| --- |
| library_name: pytorch |
| license: cc-by-nc-4.0 |
| metrics: |
| - accuracy |
| pipeline_tag: image-to-text |
| tags: |
| - OCSR |
| --- |
| |
| # COMO |
|
|
| [Online Demo](https://huggingface.co/spaces/lyuzhuoqi/COMO) \\(\cdot\\) [GitHub](https://github.com/netknowledge/COMO) \\(\cdot\\) [PyPI](https://pypi.org/project/como-ocsr/) |
|
|
| **COMO** (**C**losed-loop **O**ptical **M**olecule rec**O**gnition) is a deep learning framework for |
| Optical Chemical Structure Recognition (OCSR). It recognizes chemical structure |
| diagrams from images and predicts SMILES strings with atom-level 2D coordinates |
| and bond matrices. COMO uses Minimum Risk Training (MRT) to directly optimize |
| molecular-level, non-differentiable objectives, closing the gap between |
| token-level training and molecular-level evaluation. |
|
|
| ## Model Summary |
|
|
| - **Architecture:** Swin-B encoder β 6-layer Transformer decoder β bond MLP |
| - **Input:** 384Γ384 RGB image of a chemical structure diagram |
| - **Output:** SMILES string + atom coordinates + bond matrix |
| - **Vocabulary:** chartok_coords format (200 tokens: SMILES chars + 64 X/Y bins) |
| - **Parameters:** ~94M |
| - **Training data:** 1M PubChem + 652K USPTO (MLE) + 83K MolParser-SFT (MRT) |
| |
| ## Installation |
| |
| ```bash |
| pip install como-ocsr |
| ``` |
| |
| ## Usage |
| |
| ```python |
| import como |
| # Load pretrained checkpoints from Hugging Face |
| model = como.load_model("Keylab/COMO-Tanimoto") |
| # or download checkpoint files from HuggingFace: |
| # huggingface-cli download Keylab/COMO models/tanimoto/final.pth |
| model = como.load_model(checkpint_path="models/tanimoto/final.pth") |
|
|
| # Single image prediction |
| smiles = como.predict(model, "molecule.png") |
| print(smiles) # "CC(=O)O" |
|
|
| # Batch prediction |
| smiles_list = como.predict_batch(model, ["mol1.png", "mol2.png"]) |
|
|
| # Benchmark evaluation |
| metrics = como.evaluate(model, "benchmark/USPTO/", "benchmark/USPTO.csv") |
| print(f"Exact Match: {metrics['postprocess/exact_match_acc']:.2%}") |
| ``` |
| |
| Full documentation: [Github](https://github.com/netknowledge/COMO/tree/main/doc) |
| |
| ## Architecture |
| |
| ``` |
| Image (384Γ384) |
| β Swin-B backbone (ImageNet pretrained) |
| β 2D sinusoidal positional encoding |
| β 6-layer Transformer decoder (d=256, 8 heads) |
| β chartok_coords tokens β SMILES + coordinates |
| β Bond MLP (2-layer, GELU) β 7-class bond matrix |
| β Graph reconstruction β canonical SMILES |
| ``` |
| |
| The model outputs a molecular graph \\(G = (A, B)\\) where: |
| - \\(A = \{(l_i, x_i, y_i)\}\\) β atom SMILES labels with 2D image coordinates |
| - \\(B\\) β pairwise bond types (none, single, double, triple, aromatic, wedge, dash) |
| |
| ## Training |
| |
| ### MLE Phase |
| - **Data:** 1M PubChem SMILES (synthetic) + 652K USPTO patent molecules |
| - **Augmentation:** Indigo-rendered images with random styles, functional group |
| substitution, R-group insertion, wavy bonds, scan shadows, multilingual comments |
| - **Optimizer:** AdamW, lr=4Γ10β»β΄ (encoder & decoder), weight decay=10β»βΆ |
| - **Schedule:** 2% linear warmup β cosine decay, batch size 64/GPU |
| - **Loss:** Label-smoothed cross-entropy (Ξ΅=0.1) + bond classification CE |
| |
| ### MRT Phase |
| - **Data:** 83K real-world molecular images (MolParser-SFT) |
| - **Candidates:** N=32 per image, multinomial sampling at Ο=0.5 |
| - **Reward weights:** validity=0.1, similarity=0.5, exact match=0.4 |
| - **Sharpening:** Ξ±=1.0, loss weight Ξ»=0.1 |
| - **Schedule:** First 5 epochs MLE-only warmup, then interleaved MLE+MRT |
| |
| ## Benchmark Datasets |
| |
| Benchmark datasets are published as a HuggingFace Dataset with one config per benchmark: |
| **[Keylab/OCSR-Benchmarks](https://huggingface.co/datasets/Keylab/OCSR-Benchmarks)** |
| |
| | Config | Size | Type | |
| |--------|-----:|------| |
| | `CLEF` | 992 | Real (patents) | |
| | `JPO` | 449 | Real (patents) | |
| | `UOB` | 5,740 | Real (academic) | |
| | `USPTO` | 5,719 | Real (patents) | |
| | `USPTO-10K` | 9,999 | Real (patents) | |
| | `Staker` | 50,000 | Real | |
| | `ACS` | 331 | Real (publications) | |
| | `WildMol-10K` | 9,889 | Real (wild) | |
| | `Indigo` | 5,719 | Synthetic | |
| | `ChemDraw` | 5,719 | Synthetic | |
| |
| Each sample has three fields: `image_id` (str), `image` (PIL), `SMILES` (str). |
|
|
| ```python |
| from datasets import load_dataset |
| |
| # Load a single benchmark |
| ds = load_dataset("Keylab/OCSR-Benchmarks", name="USPTO", split="test") |
| sample = ds[0] |
| sample["image"].show() # PIL Image |
| print(sample["SMILES"]) # ground-truth SMILES |
| |
| # Iterate over all benchmarks |
| configs = ["CLEF", "JPO", "UOB", "USPTO", "USPTO-10K", |
| "Staker", "ACS", "WildMol-10K", "Indigo", "ChemDraw"] |
| for name in configs: |
| ds = load_dataset("Keylab/OCSR-Benchmarks", name=name, split="test") |
| print(f"{name}: {len(ds)} samples") |
| ``` |
|
|
| Pre-packaged `.tar.gz` archives are also available for bulk download in the |
| [COMO model repository](https://huggingface.co/Keylab/COMO/tree/main/benchmarks). |
|
|
| ## Evaluation Results |
|
|
| Exact match accuracy (%) on 10 benchmarks (COMO-Tanimoto variant): |
|
|
| | Benchmark | Images | Synthetic/Real | COMO-Tanimoto | |
| |-----------|--------|----------------|---------------| |
| | Indigo | 5,719 | Synthetic | 98.6 | |
| | ChemDraw | 5,719 | Synthetic | 96.5 | |
| | CLEF | 992 | Real (patents) | 94.8 | |
| | JPO | 449 | Real (patents) | 88.4 | |
| | UOB | 5,740 | Real (academic) | 98.0* | |
| | USPTO | 5,719 | Real (patents) | 93.4 | |
| | USPTO-10K | 9,999 | Real (patents) | 96.1 | |
| | Staker | 50,000 | Real | 87.4 | |
| | ACS | 331 | Real (publications) | 84.6 | |
| | WildMol-10K | 9,899 | Real (wild) | 77.1 | |
|
|
| *\*UOB results after tautomer standardization.* |
|
|
| See the [paper](#citation) for full comparison with MolScribe, MolParser, |
| SwinOCSR, and other baselines. |
|
|
| ## License |
|
|
| - **Code & Model Weights:** CC BY-NC 4.0 (non-commercial use only) |
| - **Benchmarks:** See original sources for applicable terms |
|
|
| ## Citation |
|
|
| ```bibtex |
| @article{lyu2026closed, |
| title={COMO: Closed-Loop Optical Molecule Recognition with Minimum Risk Training}, |
| author={Lyu, Zhuoqi and Ke, Qing}, |
| journal={arXiv preprint arXiv:2604.23546}, |
| year={2026} |
| } |
| ``` |
|
|
