MolScribe 分子结构识别模型

1、MolScribe简介

MolScribe：一种分子结构识别模型，将分子实体的图像转换为SMILES字符串。
MolScribe 项目地址

2、代码示例

pip install MolScribe
pip install huggingface_hub 

2.1 案例1

import torch
from molscribe import MolScribe
from huggingface_hub import hf_hub_download

ckpt_path = hf_hub_download('yujieq/MolScribe', 'swin_base_char_aux_1m.pth')

model = MolScribe(ckpt_path, device=torch.device('cpu'))
output = model.predict_image_file('assets/example.png', return_atoms_bonds=True, return_confidence=True)

{
    'smiles': 'Fc1ccc(-c2cc(-c3ccccc3)n(-c3ccccc3)c2)cc1',
    'molfile': '***', 
    'confidence': 0.9175,
    'atoms': [{'atom_symbol': '[Ph]', 'x': 0.5714, 'y': 0.9523, 'confidence': 0.9127}, ... ],
    'bonds': [{'bond_type': 'single', 'endpoint_atoms': [0, 1], 'confidence': 0.9999}, ... ]
}

2.2 案例2

import torch
from molscribe import MolScribe
from huggingface_hub import hf_hub_download
import os
import cv2
import numpy as np


if __name__ == "__main__":

    DATAPATH = '/mnt/d/wsl_workspace/yolo11/result/001_predict'
    PREDICT_PATH = '/mnt/d/wsl_workspace/yolo11/MolScribe/result/01'
    if not os.path.exists(PREDICT_PATH):
        os.makedirs(PREDICT_PATH)


    ckpt_path = '/mnt/d/models/MolScribe/swin_base_char_aux_1m.pth'
    model = MolScribe(ckpt_path, device=torch.device('cpu'))


    font = cv2.FONT_HERSHEY_SIMPLEX  # 字体类型
    font_scale = 0.5  # 字体大小
    color = (0, 0, 255)  # 字体颜色
    thickness = 1  # 字体粗细

    for dirpath, dirnames, filenames in os.walk(DATAPATH):
        for filename in filenames:
            inputpath = os.path.join(dirpath,filename)
            savename = os.path.join(PREDICT_PATH,filename)
            image = cv2.imread(inputpath)

            large_matrix = np.full((300,500,3),255,dtype=np.uint8)
            start_row = (large_matrix.shape[0] - image.shape[0])//2
            start_col = (large_matrix.shape[1] - image.shape[1])//2
            large_matrix[start_row:start_row + image.shape[0], 
                         start_col:start_col + image.shape[1],:] = image
            
            output = model.predict_image_file(inputpath, return_atoms_bonds=True, return_confidence=True)
            print(filename,' : ',output['smiles'])
            cv2.putText(large_matrix, output['smiles'], (0, 290), font, font_scale, color, thickness)
            cv2.imwrite(savename,large_matrix)