Skip to content

Pic2SMILES

Optical chemical structure recognition (OCSR): turn a picture of a 2D chemical structure into canonical SMILES.

Try it in your browser -- the recognizer runs on your own device, as a static page with no server behind it. Reports and documentation: https://pic2smiles.pages.dev/.

This repository trains MolMini, a small image-to-SMILES encoder-decoder, and benchmarks it honestly against the alternatives. The model is deliberately laptop-sized: training happens wherever a GPU is available, but inference is designed to run on a 16 GB Apple Silicon MacBook with nothing but PyTorch and RDKit installed.

What is here

Layer Path Purpose
chemistry src/pic2smiles/chem/ canonicalization, atom-level SMILES tokenizer, OCSR metrics, corpus curation with scope + leakage guards
data src/pic2smiles/data/ randomized RDKit depiction renderer, image corruption pipeline, on-the-fly dataset, frozen eval manifests
model src/pic2smiles/models/ MolMini: conv-stem + 2D-RoPE transformer encoder, cross-attending causal SMILES decoder
training src/pic2smiles/train/ step-based trainer, WSD schedule, weight EMA, resumable checkpoints
inference src/pic2smiles/infer/ KV-cache greedy/beam decoding, MolMiniPredictor
evaluation src/pic2smiles/evaluation/ recognizer contract, benchmark harness, Markdown report

Quick start

pip install -e ".[dev]"
pic2smiles env

Curate a molecule corpus (any newline-delimited SMILES file works):

pic2smiles corpus --source chembl_smiles.txt --output data/corpus/chembl --holdout data/corpus/decimer_holdout.smi

Freeze an evaluation set and train:

pic2smiles eval-set --smiles data/corpus/chembl/test.smi --output data/eval/chembl_test --count 1000
pic2smiles train --corpus data/corpus/chembl --output runs/molmini --preset base --steps 60000

Predict and benchmark:

pic2smiles predict --checkpoint runs/molmini/best.pt path/to/structure.png
pic2smiles predict --checkpoint runs/molmini/best.pt --screenshot path/to/screenshot.png
pic2smiles benchmark --manifest data/eval/chembl_test/manifest.csv --dataset chembl_test \
  --models molmini oracle image_hash_nn --checkpoint runs/molmini/best.pt
pic2smiles report --results reports/results --output reports/benchmark_report.md

Pass --screenshot when the image is a screenshot, photo or scan rather than a render: it flattens transparency, inverts dark-mode captures, trims margins and squares the picture before inference, and reports which of those it applied. Leave it off for renders -- measured on 300 frozen render images it costs about four points of exact match (0.877 to 0.833), because trimming a render's designed margin changes the scale the model was trained on.

MolMini, in a browser

Export an inference checkpoint and serve the local web UI: paste a structure screenshot, get a copyable SMILES with a depiction and computed properties.

pic2smiles export --checkpoint checkpoints/molmini_base_last.pt --output molmini.pt --half
pic2smiles serve --open

serve finds molmini.pt (or checkpoints/molmini*.pt) on its own; point it elsewhere with --checkpoint or MOLMINI_CHECKPOINT. Other flags: --host, --port, --device, --beam-size, --no-preload.

  • Paste (Cmd+V), drag-drop or pick a file. Pasting SMILES text skips the model and goes straight to the depiction and properties.
  • Screenshots are normalized into the training distribution before inference: transparency flattened, dark-mode captures inverted, uniform margins trimmed, and padded to a square. That last step matters more than it sounds -- training depictions are square, the predictor resizes without preserving aspect ratio, and a wide screenshot fed in raw decodes into a confidently wrong molecule.
  • The predicted SMILES is editable and everything recomputes as you type, so a near-miss can be corrected in place.
  • Source screenshot and RDKit re-render sit side by side. The decode confidence is shown but labelled as what it is: a mean per-token probability, not a correctness estimate. The visual comparison is the check that matters.
  • A prediction outside the model card's declared scope (heavy-atom count, elements, multi-fragment) is flagged rather than reported as a normal result.

--half halves the exported file (381 MB to 48 MB); it does not change results or speed, because the predictor loads the weights back into a float32 model.

Measured on an M3 (CPU, greedy): 1.9 s to load, ~110 ms per small molecule and ~200 ms for a 37-heavy-atom drug. CPU is the default because it beats MPS at this model size -- autoregressive decoding is many small kernels, where MPS dispatch overhead dominates.

Model

MolMini presets, all trained the same way:

preset parameters input encoder decoder
tiny 3.6 M 256 px conv stem + 3 blocks, d=192 3 blocks, d=192
small 12.1 M 320 px conv stem + 4 blocks, d=288 5 blocks, d=288
base 23.8 M 384 px conv stem + 4 blocks, d=384 6 blocks, d=384

The encoder is a hybrid on purpose. Convolutions carry 384x384 pixels down to a 24x24 grid cheaply, because early OCSR features (strokes, junctions, glyphs) are local; global attention then runs only on the 576 surviving tokens, where it is actually needed to relate the two halves of a ring closure. Blocks use RMSNorm, SwiGLU, QK-norm and rotary positions -- 2D over the image grid, 1D over the token sequence.

Data

Training images are generated on the fly rather than frozen to disk. Rendering one depiction costs about 2 ms of CPU, far less than one GPU step, so a fixed image set would only give the model something to memorize. Every render randomizes layout engine, rotation, bond width, font, atom palette, condensed abbreviations (Ph, OMe) and drawing style; the image is then corrupted with stroke morphology, blur, downscaling, JPEG artifacts, perspective warp, paper texture and speckle.

Evaluation images are the opposite: written once to disk with a fixed seed, in three tiers -- clean (textbook depiction), varied (publication style diversity) and degraded (scan and photo conditions).

Honesty rules this repo follows

  • Every rate is computed over molecules the model never saw during training; splits are disjoint by canonical SMILES and asserted, not assumed.
  • Held-out real datasets are subtracted from the training corpus, and the number of molecules removed is recorded in corpus_meta.json.
  • A system that cannot run is recorded as skipped with the reason and is excluded from every rate, rather than scoring zero.
  • Reports state the declared molecular scope. Numbers do not transfer outside it.

Documentation