repository·main·Indexed 43 days ago
https://github.com/thomas0809/molscribeAn image-to-graph model that translates molecular images into chemical structures, including SMILES, Molfiles, and atom/bond graphs. It provides a Python API and CLI for predicting molecular structures, tools for evaluating prediction accuracy against gold standards, and utilities for canonicalizing SMILES and visualizing predicted graphs overlaid on original images.
You can install MolScribe using one of two methods:
Option 1: via pip
pip install MolScribeOption 2: from source Clone the repository and run the setup script:
git clone git@github.com:thomas0809/MolScribe.git
cd MolScribe
python setup.py installpip install MolScribeFor reproducing experiments or training the model, follow these steps:
1. Install dependencies:
pip install -r requirements.txt2. Prepare Data:
Download datasets (USPTO, PubChem, Synthetic, Realistic, or Perturbed) and place them in the data/ directory.
3. Download Checkpoints:
Create a ckpts directory and download the required weights:
mkdir -p ckpts
wget -P ckpts https://huggingface.co/yujieq/MolScribe/resolve/main/swin_base_char_aux_1m680k.pth4. Run Training/Evaluation Scripts:
bash scripts/eval_uspto_joint_chartok_1m680k.sh (Uses 1 GPU, batch size 64 by default).bash scripts/train_uspto_joint_chartok_1m680k.sh (Uses 4 GPUs, batch size 256 by default).To perform molecular structure prediction from an image, use the MolScribe class. You can initialize the model with a checkpoint path and a target device (e.g., cpu or cuda).
Use predict_image_file to process a single image. Setting return_atoms_bonds=True ensures the prediction dictionary contains both the SMILES string and the MolBlock.
To process a list of images at once, use predict_image_files.
from molscribe import MolScribe
import torch
MODEL_PATH = '../ckpts/swin_base_char_aux_1m680k.pth'
IMAGE_PATH = '../assets/example.png'
device = torch.device('cpu')
model = MolScribe(MODEL_PATH, device)
# Single image inference
prediction = model.predict_image_file(IMAGE_PATH, return_atoms_bonds=True)
smiles = prediction['smiles']
molfile = prediction['molfile']
# Multiple images inference
more_images = ['../assets/example1.png', '../assets/example2.png']
predictions = model.predict_image_files(more_images)You can troubleshoot MolScribe predictions by comparing ground truth SMILES with predicted SMILES (SMILES), post-processed SMILES (post_SMILES), and graph-derived SMILES (graph_SMILES).
To visualize where the model is failing, you can overlay the predicted molecular graph (nodes and edges) onto the original image. The graph components are typically stored in the prediction CSV as:
node_coords: Normalized coordinates of atoms.node_symbols: Atomic symbols.edges: Adjacency matrix representing bond types.When visualizing, remember that coordinates in the CSV are normalized; you must scale them by the image width and height to overlay them correctly.
# Example logic for scaling normalized coordinates to image dimensions
h, w, _ = img.shape
x, y = coords[:,0]*w, coords[:,1]*h
# Plotting nodes and symbols
plt.scatter(x, y, color='r', marker='o')
for i in range(len(symbols)):
plt.text(x[i], y[i], symbols[i], color='blue')
# Plotting edges (bonds)
for i in range(len(x)):
for j in range(len(x)):
if edges[i][j] != 0:
# Logic for different bond types (e.g., 5, 6, 1, etc.)
plt.plot([x[i], x[j]], [y[i], y[j]], color)To use MolScribe in your Python code, download a checkpoint from the HuggingFace Hub and use the MolScribe class. The predict_image_file method allows you to extract chemical structures from images.
Output Format: The method returns a dictionary containing:
smiles: The predicted SMILES string.molfile: The predicted Molfile string.confidence: A float representing prediction confidence.atoms: A list of dictionaries, each containing atom_symbol, x, y coordinates, and confidence.bonds: A list of dictionaries, each containing bond_type, endpoint_atoms (indices), and confidence.import torch
from molscribe import MolScribe
from huggingface_hub import hf_hub_download
# Download checkpoint
ckpt_path = hf_hub_download('yujieq/MolScribe', 'swin_base_char_aux_1m.pth')
# Initialize model
model = MolScribe(ckpt_path, device=torch.device('cpu'))
# Predict
output = model.predict_image_file('assets/example.png', return_atoms_bonds=True, return_confidence=True)To evaluate the model's performance, use the evaluate.py script. This script compares predicted SMILES against a gold standard file.
Requirements for prediction files: Predictions should be saved in a CSV file with:
image_id column (must match the gold file).SMILES column (or another column name specified via --pred_field).Metrics provided:
canon_smiles: Exact matching accuracy (main metric).graph: Graph exact matching accuracy (ignores tetrahedral chirality).chiral: Exact matching accuracy on chiral molecules.python evaluate.py \
--gold_file data/real/acs.csv \
--pred_file output/uspto/swin_base_char_aux_1m680k/prediction_acs.csv \
--pred_field post_SMILESMolScribe provides a built-in method draw_prediction to visualize the model's output overlaid on the original image. This is useful for verifying atom and bond assignments.
Note: When viewing predictions via RDKit (Chem.MolFromMolBlock), superatoms may not be displayed correctly. For full structural fidelity including superatoms, it is recommended to copy the molfile output and paste it into ChemDraw, ensuring the leading blank line is included.
import cv2
import matplotlib.pyplot as plt
image = cv2.imread(IMAGE_PATH)
model.draw_prediction(prediction, image, notebook=True)You can run predictions directly from the command line using the predict.py script. You must provide the path to a model checkpoint and the path to the image file.
python predict.py --model_path ckpts/swin_base_char_aux_1m680k.pth --image_path assets/example.pngPerforms inference on a list of image files.
Parameters:
image_paths (list of str): A list of paths to the input images.Returns:
list[dict]: A list of prediction dictionaries, one for each input image.To ensure fair comparison between ground truth and predicted SMILES, use canonicalize_smiles from molscribe.evaluate. This function helps handle variations in SMILES representation by producing a canonical form. It supports options to ignore stereochemistry or cis/trans isomerism.
Arguments:
smiles: The SMILES string to canonicalize.ignore_cistrans (bool): Whether to ignore cis/trans isomerism.ignore_chiral (bool): Whether to ignore chirality.Returns: A tuple where the first element is the canonicalized SMILES string.
from molscribe.evaluate import canonicalize_smiles
gold_smiles, _ = canonicalize_smiles(row['SMILES'], ignore_cistrans=True, ignore_chiral=False)Performs inference on a single image file.
Parameters:
image_path (str): Path to the input image.return_atoms_bonds (bool): If True, the returned dictionary includes atom and bond information (as a MolBlock).Returns:
dict: A dictionary containing:'smiles' (str): The predicted SMILES string.'molfile' (str, optional): The predicted MolBlock (if return_atoms_bonds=True).The _replace_functional_group function is used to process SMILES strings that may contain placeholders or specific representations for functional groups, converting them into standard SMILES notation.
from molscribe.chemistry import _replace_functional_group
# pred_smiles is the raw prediction
pred_smiles, mappings = _replace_functional_group(pred_smiles)