from rdkit import Chem
from rdkit.Chem.Draw import rdMolDraw2D
from IPython.display import SVG
from io import BytesIO
from PIL import Image
from cairosvg import svg2png
import argparsedef generate_image(mol, highlight_atoms, highlight_bonds, atomColors, bondColors, radii, size, output, isNumber=False):print(highlight_atoms)print(highlight_bonds)print(atomColors)print(bondColors)image_data = BytesIO()view = rdMolDraw2D.MolDraw2DSVG(size[0], size[1])tm = rdMolDraw2D.PrepareMolForDrawing(mol)option = view.drawOptions()if isNumber:for atom in mol.GetAtoms():option.atomLabels[atom.GetIdx()] = atom.GetSymbol() + str(atom.GetIdx() + 1)view.DrawMolecule(tm, highlightAtoms=highlight_atoms, highlightBonds=highlight_bonds,highlightAtomColors=atomColors, highlightBondColors=bondColors, highlightAtomRadii=radii)view.FinishDrawing()svg = view.GetDrawingText()SVG(svg.replace('svg:', ''))svg2png(bytestring=svg, write_to=output)img = Image.open(output)img.save(image_data, format='PNG')return image_datadef main():parser = argparse.ArgumentParser()parser.add_argument("-i", type=str, required=True)parser.add_argument("-o", type=str, required=True)parser.add_argument("-a", type=str, action="append")parser.add_argument("-b", type=str, action="append")args = parser.parse_args()atom_hilights = {}bond_hilights = {}radii = {}with open(args.i, "r") as f:string_mol = f.read()mol = Chem.MolFromMolBlock(string_mol)if args.a:for atom in args.a:print(atom)datas = atom.split(",")atom_hilights[int(datas[0])] = (int(datas[1]), int(datas[2]), int(datas[3]))radii[int(datas[0])] = 0.20if args.b:for bond in args.b:datas = bond.split(",")bond_hilights[int(datas[0])] = (int(datas[1]), int(datas[2]), int(datas[3]))generate_image(mol,list(atom_hilights.keys()), list(bond_hilights.keys()), atom_hilights, bond_hilights, radii, (400, 200), args.o, False)if __name__ == "__main__":main()