Blame for README.md

---
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}
}
```