decimer 2.7.0__tar.gz → 2.8.0__tar.gz

{decimer-2.7.0 → decimer-2.8.0}/DECIMER/Efficient_Net_encoder.py

@@ -1,8 +1,6 @@
 # EfficientNet-V2 config
 import tensorflow as tf
 
-import DECIMER.efficientnetv2
-from DECIMER.efficientnetv2 import effnetv2_configs
 from DECIMER.efficientnetv2 import effnetv2_model
 
 BATCH_SIZE_DEBUG = 2

decimer-2.8.0/DECIMER/__init__.py

@@ -0,0 +1,31 @@
+# -*- coding: utf-8 -*-
+"""
+DECIMER V2.7.2 Python Package
+=============================
+
+Deep lEarning for Chemical ImagE Recognition (DECIMER) project
+was launched to address the Optical Chemical Structure Recognition (OCSR)
+problem using deep learning based methods,
+providing an automated open-source software solution.
+
+Typical usage example::
+
+    from decimer import predict_SMILES
+
+    # Chemical depiction to SMILES translation
+    image_path = "path/to/imagefile"
+    SMILES = predict_SMILES(image_path)
+    print(SMILES)
+
+For comments, bug reports, or feature ideas,
+please raise an issue on the GitHub repository.
+"""
+
+__version__ = "2.7.2"
+
+__all__ = [
+    "DECIMER",
+]
+
+
+from .decimer import predict_SMILES

decimer-2.8.0/DECIMER/config.py

@@ -0,0 +1,148 @@
+import tensorflow as tf
+
+import DECIMER.Efficient_Net_encoder as Efficient_Net_encoder
+import DECIMER.Transformer_decoder as Transformer_decoder
+
+TARGET_DTYPE = tf.float32
+
+
+class Config:
+    """Configuration class."""
+
+    def __init__(
+        self,
+    ):
+        self.encoder_config = {}
+        self.transformer_config = {}
+        self.lr_config = {}
+
+    def initialize_encoder_config(
+        self,
+        image_embedding_dim,
+        preprocessing_fn,
+        backbone_fn,
+        image_shape,
+        do_permute=False,
+        pretrained_weights=None,
+    ):
+        """This functions initializes the Efficient-Net V2 encoder with user
+        defined configurations.
+
+        Args:
+            image_embedding_dim (int): Embedding dimension of the input image
+            preprocessing_fn (method): Efficient Net preprocessing function for input image
+            backbone_fn (method): Calls Efficient-Net V2 as backbone for encoder
+            image_shape (int): Shape of the input image
+            do_permute (bool, optional): . Defaults to False.
+            pretrained_weights (keras weights, optional): Use pretrained efficient net weights or not. Defaults to None.
+        """
+        self.encoder_config = dict(
+            image_embedding_dim=image_embedding_dim,
+            preprocessing_fn=preprocessing_fn,
+            backbone_fn=backbone_fn,
+            image_shape=image_shape,
+            do_permute=do_permute,
+            pretrained_weights=pretrained_weights,
+        )
+
+    def initialize_transformer_config(
+        self,
+        vocab_len,
+        max_len,
+        n_transformer_layers,
+        transformer_d_dff,
+        transformer_n_heads,
+        image_embedding_dim,
+        rate=0.1,
+    ):
+        """This functions initializes the Transformer model as decoder with
+        user defined configurations.
+
+        Args:
+            vocab_len (int): Total number of words in the input vocabulary
+            max_len (int): Maximum length of the string found on the training dataset
+            n_transformer_layers (int): Number of layers present in the transformer model
+            transformer_d_dff (int): Transformer feed forward upwards projection size
+            transformer_n_heads (int): Number of heads present in the transformer model
+            image_embedding_dim (int): Total number of dimension the image gets embedded
+            dropout_rate (float, optional): Fraction of the input units to drop. Defaults to 0.1.
+        """
+        self.transformer_config = dict(
+            num_layers=n_transformer_layers,
+            d_model=image_embedding_dim,
+            num_heads=transformer_n_heads,
+            dff=transformer_d_dff,
+            target_vocab_size=vocab_len,
+            max_len=max_len,
+            rate=0.1,
+        )
+
+    def initialize_lr_config(self, warm_steps, n_epochs):
+        """This function sets the configuration to initialize learning rate.
+
+        Args:
+            warm_steps (int): Number of steps The learning rate is increased
+            n_epochs (int): Number of epochs
+        """
+        self.lr_config = dict(
+            warm_steps=warm_steps,
+            n_epochs=n_epochs,
+        )
+
+
+class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
+    """Custom schedule for learning rate used during training.
+
+    Args:
+        tf (_type_): keras learning rate schedule
+    """
+
+    def __init__(self, d_model, warmup_steps=4000):
+        super(CustomSchedule, self).__init__()
+
+        self.d_model = d_model
+        self.d_model = tf.cast(self.d_model, tf.float32)
+
+        self.warmup_steps = warmup_steps
+
+    def __call__(self, step):
+        arg1 = tf.math.rsqrt(step)
+        arg2 = step * (self.warmup_steps**-1.5)
+
+        return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)
+
+
+def prepare_models(encoder_config, transformer_config, replica_batch_size, verbose=0):
+    """This function is used to initiate the Encoder and the Transformer with
+    appropriate configs set by the user. After initiating the models this
+    function returns the Encoder,Transformer and the optimizer.
+
+    Args:
+        encoder_config ([type]): Encoder configuration set by user in the config class.
+        transformer_config ([type]): Transformer configuration set by user in the config class.
+        replica_batch_size ([type]): Per replica batch size set by user(during distributed training).
+        verbose (int, optional): Defaults to 0.
+
+    Returns:
+        [type]: Optimizer, Encoder model and the Transformer
+    """
+
+    # Instantiate an optimizer
+    optimizer = tf.keras.optimizers.Adam(learning_rate=0.00051)
+
+    # Instantiate the encoder model
+
+    encoder = Efficient_Net_encoder.Encoder(**encoder_config)
+
+    # Instantiate the decoder model
+    transformer = Transformer_decoder.Decoder(**transformer_config)
+
+    # Show the model architectures and plot the learning rate
+    if verbose:
+        print("\nEncoder model\n")
+        print(encoder.summary())
+
+        print("\nTransformer model\n")
+        print(transformer.summary())
+
+    return optimizer, encoder, transformer

{decimer-2.7.0 → decimer-2.8.0}/DECIMER/decimer.py

@@ -5,10 +5,11 @@ import sys
 from typing import List
 from typing import Tuple
 
+import numpy as np
 import pystow
 import tensorflow as tf
 
-import DECIMER.config as config
+import DECIMER.pre_process as pre_process
 import DECIMER.utils as utils
 
 # Silence tensorflow model loading warnings.
@@ -17,6 +18,46 @@ logging.getLogger("absl").setLevel("ERROR")
 # Silence tensorflow errors - not recommended if your model is not working properly.
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
 
+
+# ============================================================================
+# Keras 2 to Keras 3 Compatibility Layer
+# ============================================================================
+class KerasCompatibilityUnpickler(pickle.Unpickler):
+    """Custom unpickler to handle Keras 2.x tokenizers in Keras 3 environment.
+
+    This handles the module path changes between Keras 2 and Keras 3:
+    - keras.preprocessing.text -> tensorflow.keras.preprocessing.text
+    """
+
+    def find_class(self, module, name):
+        # Redirect old Keras 2.x imports to TensorFlow Keras
+        if module.startswith("keras."):
+            module = module.replace("keras.", "tensorflow.keras.")
+        return super().find_class(module, name)
+
+
+def load_tokenizer(tokenizer_path: str) -> object:
+    """Load tokenizer with Keras 2/3 compatibility.
+
+    Args:
+        tokenizer_path (str): Path to the pickled tokenizer file
+
+    Returns:
+        tokenizer: Loaded tokenizer object
+    """
+    try:
+        # Try standard pickle load first (for Keras 3 tokenizers)
+        with open(tokenizer_path, "rb") as f:
+            return pickle.load(f)
+    except ModuleNotFoundError as e:
+        if "keras.preprocessing" in str(e):
+            # Use compatibility unpickler for Keras 2.x tokenizers
+            with open(tokenizer_path, "rb") as f:
+                return KerasCompatibilityUnpickler(f).load()
+        else:
+            raise
+
+
 # Set the absolute path
 HERE = os.path.dirname(os.path.abspath(__file__))
 
@@ -58,7 +99,7 @@ def get_models(model_urls: dict):
     tokenizer_path = os.path.join(
         model_paths["DECIMER"], "assets", "tokenizer_SMILES.pkl"
     )
-    tokenizer = pickle.load(open(tokenizer_path, "rb"))
+    tokenizer = load_tokenizer(tokenizer_path)
 
     # Load DECIMER models
     DECIMER_V2 = tf.saved_model.load(model_paths["DECIMER"])
@@ -122,19 +163,19 @@ def detokenize_output_add_confidence(
 
 
 def predict_SMILES(
-    image_path: str, confidence: bool = False, hand_drawn: bool = False
+    image_input: [str, np.ndarray], confidence: bool = False, hand_drawn: bool = False
 ) -> str:
     """Predicts SMILES representation of a molecule depicted in the given image.
 
     Args:
-        image_path (str): Path of chemical structure depiction image
-        confidence (bool): Flag to indicate whether to return confidence values along with SMILES prediction
-        hand_drawn (bool): Flag to indicate whether the molecule in the image is hand-drawn
+        image_input (str or np.ndarray): Path of chemical structure depiction image or a numpy array representing the image.
+        confidence (bool): Flag to indicate whether to return confidence values along with SMILES prediction.
+        hand_drawn (bool): Flag to indicate whether the molecule in the image is hand-drawn.
 
     Returns:
-        str: SMILES representation of the molecule in the input image, optionally with confidence values
+        str: SMILES representation of the molecule in the input image, optionally with confidence values.
     """
-    chemical_structure = config.decode_image(image_path)
+    chemical_structure = pre_process.decode_image(image_input)
 
     model = DECIMER_Hand_drawn if hand_drawn else DECIMER_V2
     predicted_tokens, confidence_values = model(tf.constant(chemical_structure))

decimer-2.8.0/DECIMER/pre_process.py

@@ -0,0 +1,243 @@
+import io
+import cv2
+import efficientnet.tfkeras as efn
+import tensorflow as tf
+import numpy as np
+
+from PIL import Image
+from PIL import ImageEnhance
+from pillow_heif import register_heif_opener
+from typing import Union
+
+
+def resize_byratio(image):
+    """This function takes a Pillow Image object and will resize the image by
+    512 x 512 To upscale or to downscale the image LANCZOS resampling method is
+    used.
+
+    with the new pillow version the antialias is turned on when using LANCZOS.
+    Args: PIL.Image
+    Returns: PIL.Image
+    """
+    maxwidth = 512
+    ratio = maxwidth / max(image.width, image.height)
+    new_size = int((float(image.width) * ratio)), int((float(image.height) * ratio))
+    resized_image = image.resize(new_size, resample=Image.LANCZOS)
+    return resized_image
+
+
+def central_square_image(image):
+    """This function takes a Pillow Image object and will add white padding so
+    that the image has a square shape with the width/height of the longest side
+    of the original image.
+
+    Args: PIL.Image
+    Returns: PIL.Image
+    """
+    max_wh = int(1.2 * max(image.size))
+    # If the new image is smaller than 299x299, then let's paste it into an empty image
+    # of that size instead of distorting it later while resizing.
+    if max_wh < 512:
+        max_wh = 512
+    new_im = Image.new(image.mode, (max_wh, max_wh), "white")
+    paste_pos = (
+        int((new_im.size[0] - image.size[0]) / 2),
+        int((new_im.size[1] - image.size[1]) / 2),
+    )
+    new_im.paste(image, paste_pos)
+    return new_im
+
+
+def delete_empty_borders(image):
+    """This function takes a Pillow Image object, converts it to grayscale and
+    deletes white space at the borders.
+
+    Args: PIL.Image
+    Returns: PIL.Image
+    """
+    image = np.asarray(image.convert("L"))
+    mask = image > 200
+    rows = np.flatnonzero((~mask).sum(axis=1))
+    cols = np.flatnonzero((~mask).sum(axis=0))
+    crop = image[rows.min() : rows.max() + 1, cols.min() : cols.max() + 1]
+    return Image.fromarray(crop)
+
+
+def PIL_im_to_BytesIO(image):
+    """
+    Convert pillow image to io.BytesIO object
+    Args: PIL.Image
+    Returns: io.BytesIO object with the image data
+    """
+    output = io.BytesIO()
+    image.save(output, format="PNG")
+    return output
+
+
+def HEIF_to_pillow(image_path: str):
+    """
+    Converts Apple's HEIF format to useful pillow object
+    Returns: image_path (str): path of input image
+    Returns: PIL.Image
+    """
+    register_heif_opener()
+
+    heif_file = Image.open(image_path).convert("RGBA")
+    return heif_file
+
+
+def remove_transparent(image: Union[str, np.ndarray]) -> Image.Image:
+    """
+    Removes the transparent layer from a PNG image with an alpha channel.
+
+    Args:
+        image (Union[str, np.ndarray]): Path of the input image or a numpy array representing the image.
+
+    Returns:
+        PIL.Image.Image: The image with transparency removed.
+    """
+
+    def process_image(png: Image.Image) -> Image.Image:
+        """
+        Helper function to remove transparency from a single image.
+
+        Args:
+            png (PIL.Image.Image): The input PIL image with transparency.
+
+        Returns:
+            PIL.Image.Image: The image with transparency removed.
+        """
+        background = Image.new("RGBA", png.size, (255, 255, 255))
+        alpha_composite = Image.alpha_composite(background, png)
+        return alpha_composite
+
+    def handle_image_path(image_path: str) -> Image.Image:
+        """
+        Helper function to handle image paths.
+
+        Args:
+            image_path (str): The path to the input image.
+
+        Returns:
+            PIL.Image.Image: The image with transparency removed.
+        """
+        try:
+            png = Image.open(image_path).convert("RGBA")
+        except Exception as e:
+            if type(e).__name__ == "UnidentifiedImageError":
+                png = HEIF_to_pillow(image_path)
+            else:
+                print(e)
+                raise Exception
+        return process_image(png)
+
+    def handle_numpy_array(array: np.ndarray) -> Image.Image:
+        """
+        Helper function to handle a numpy array.
+
+        Args:
+            array (np.ndarray): The numpy array representing the image.
+
+        Returns:
+            PIL.Image.Image: The image with transparency removed.
+        """
+        png = Image.fromarray(array).convert("RGBA")
+        return process_image(png)
+
+    # Check if input is a numpy array
+    if isinstance(image, np.ndarray):
+        return handle_numpy_array(array=image)
+
+    return handle_image_path(image_path=image)
+
+
+def get_bnw_image(image):
+    """
+    converts images to black and white
+    Args: PIL.Image
+    Returns: PIL.Image
+    """
+
+    im_np = np.asarray(image)
+    grayscale = cv2.cvtColor(im_np, cv2.COLOR_BGR2GRAY)
+    # (thresh, im_bw) = cv2.threshold(grayscale, 128, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
+    im_pil = Image.fromarray(grayscale)
+    enhancer = ImageEnhance.Contrast(im_pil)
+    im_output = enhancer.enhance(1.8)
+    return im_output
+
+
+def increase_contrast(image):
+    """This function increases the contrast of an image input.
+
+    Args: PIL.Image
+    Returns: PIL.Image
+    """
+
+    # Get brightness range - cast to int to avoid uint8 overflow
+    min_val = int(np.min(image))
+    max_val = int(np.max(image))
+
+    # Use LUT to convert image values
+    LUT = np.zeros(256, dtype=np.uint8)
+    LUT[min_val : max_val + 1] = np.linspace(
+        start=0, stop=255, num=(max_val - min_val) + 1, endpoint=True, dtype=np.uint8
+    )
+
+    # Apply LUT and return image
+    return Image.fromarray(LUT[image])
+
+
+def get_resize(image):
+    """This function used to decide how to resize a given image without losing
+    much information.
+
+    Args: PIL.Image
+    Returns: PIL.Image
+    """
+    width, height = image.size
+
+    # Checks the image size and resizes using the LANCOS image resampling
+    if (width == height) and (width < 512):
+        image = image.resize((512, 512), resample=Image.LANCZOS)
+    elif width >= 512 or height >= 512:
+        return image
+    else:
+        image = resize_byratio(image)
+
+    return image
+
+
+def increase_brightness(image):
+    """This function adjusts the brightness of the given image.
+
+    Args: PIL.Image
+    Returns: PIL.Image
+    """
+    enhancer = ImageEnhance.Brightness(image)
+    image = enhancer.enhance(1.6)
+    return image
+
+
+def decode_image(image_path: Union[str, np.ndarray]):
+    """Loads an image and preprocesses the input image in several steps to get
+    the image ready for DECIMER input.
+
+    Args:
+        image_path (Union[str, np.ndarray]): path of input image or numpy array representing the image.
+
+    Returns:
+        Processed image
+    """
+    img = remove_transparent(image_path)
+    img = increase_contrast(img)
+    img = get_bnw_image(img)
+    img = get_resize(img)
+    img = delete_empty_borders(img)
+    img = central_square_image(img)
+    img = increase_brightness(img)
+    img = PIL_im_to_BytesIO(img)
+    img = tf.image.decode_png(img.getvalue(), channels=3)
+    img = tf.image.resize(img, (512, 512), method="gaussian", antialias=True)
+    img = efn.preprocess_input(img)
+    return img

decimer-2.8.0/DECIMER/utils.py

@@ -0,0 +1,143 @@
+import os
+import re
+import shutil
+
+import pystow
+import zipfile
+from pathlib import Path
+
+pattern = "R([0-9]*)|X([0-9]*)|Y([0-9]*)|Z([0-9]*)"
+add_space_re = r"^(\W+)|(\W+)$"
+
+
+def split_and_modify_atoms(SMILES):
+    splitted_SMILES = list(SMILES)
+    modified_SMILES = re.sub(r"\s+(?=[a-z])", "", " ".join(map(str, splitted_SMILES)))
+    return modified_SMILES
+
+
+def replacer(match):
+    part = match.group(0)
+    text = (
+        part.replace("1", "!")
+        .replace("2", "$")
+        .replace("3", "^")
+        .replace("4", "<")
+        .replace("5", ">")
+        .replace("6", "?")
+        .replace("7", "£")
+        .replace("8", "¢")
+        .replace("9", "€")
+        .replace("0", "§")
+    )
+    return text
+
+
+def add_space(match_):
+    if match_.group(1) is not None:
+        return "{} ".format(match_.group(1))
+    else:
+        return " {}".format(match_.group(2))
+
+
+def encoder(SMILES):
+    replaced_SMILES = re.sub(pattern, replacer, SMILES)
+    splitted_SMILES = split_and_modify_atoms(replaced_SMILES)
+    modified_SMILES = " ".join(
+        [re.sub(add_space_re, add_space, word) for word in splitted_SMILES.split()]
+    )
+    return modified_SMILES
+
+
+def decoder(predictions):
+    modified = (
+        predictions.replace("!", "1")
+        .replace("$", "2")
+        .replace("^", "3")
+        .replace("<", "4")
+        .replace(">", "5")
+        .replace("?", "6")
+        .replace("£", "7")
+        .replace("¢", "8")
+        .replace("€", "9")
+        .replace("§", "0")
+    )
+    return modified
+
+
+# Downloads the model and unzips the file downloaded, if the model is not present on the working directory.
+def download_trained_weights(model_url: str, model_path: str, verbose=1):
+    """This function downloads the trained models and tokenizers to a default
+    location. After downloading the zipped file the function unzips the file
+    automatically. If the model exists on the default location this function
+    will not work.
+
+    Args:
+        model_url (str): trained model url for downloading.
+        model_path (str): model default path to download.
+
+    Returns:
+        path (str): downloaded model.
+    """
+    # Download trained models
+    if verbose > 0:
+        print("Downloading trained model to " + str(model_path))
+    model_path = pystow.ensure("DECIMER-V2", url=model_url)
+    if verbose > 0:
+        print(model_path)
+        print("... done downloading trained model!")
+
+    with zipfile.ZipFile(model_path.as_posix(), "r") as zip_ref:
+        zip_ref.extractall(model_path.parent.as_posix())
+
+    # Delete zipfile after downloading
+    if Path(model_path).exists():
+        Path(model_path).unlink()
+
+
+def ensure_models(default_path: str, model_urls: dict) -> dict:
+    """Function to ensure models are present locally.
+
+    Convenient function to ensure model downloads before usage.
+    Models are re-downloaded if the URL changes (e.g., new Zenodo record).
+
+    Args:
+        default_path (str): Default path for model data
+        model_urls (dict): Dictionary containing model names as keys and their corresponding URLs as values
+
+    Returns:
+        dict: A dictionary containing model names as keys and their local paths as values
+    """
+    model_paths = {}
+
+    for model_name, model_url in model_urls.items():
+        model_path = os.path.join(default_path, f"{model_name}_model")
+        saved_model_file = os.path.join(model_path, "saved_model.pb")
+        version_file = os.path.join(model_path, ".model_url")
+
+        # Check if model needs to be downloaded:
+        # 1. saved_model.pb doesn't exist, OR
+        # 2. The URL has changed (model was updated)
+        needs_download = not os.path.exists(saved_model_file)
+
+        if not needs_download and os.path.exists(version_file):
+            with open(version_file, "r") as f:
+                cached_url = f.read().strip()
+            if cached_url != model_url:
+                needs_download = True
+                print(f"Model {model_name} has been updated, re-downloading...")
+
+        if needs_download:
+            # Clean up incomplete/corrupted model directory if it exists
+            if os.path.exists(model_path):
+                shutil.rmtree(model_path)
+            download_trained_weights(model_url, default_path)
+
+            # Store the URL used for this download
+            os.makedirs(model_path, exist_ok=True)
+            with open(version_file, "w") as f:
+                f.write(model_url)
+
+        model_paths[model_name] = model_path
+
+    return model_paths