import os
import sys
from .xvfbwrapper import Xvfb
[docs]class ImageGenerator:
"""
A base class to generate 2D structure image file from SMILES
"""
[docs] def __init__(self):
self._supported_formats = {}
[docs] def generate(self, smiles, image, format=None):
"""
Generating 2D image file from SMILES string
:type smiles: string
:param smiles: SMILES string representing molecule structure.
:type image: string
:param image: The output image file name of 2D structure.
:type format: string
:param format: output format of the image file. By default, the output
format will be deduced from the output image file name.
"""
if format is None:
root, ext = os.path.splitext(image)
if not ext:
raise ValueError("can not determine image format from suffix")
format = ext[1:].lower()
else:
format = format.lower()
if format not in list(self._supported_formats):
raise ValueError("format (%s) is not supported" % format)
image_generator = self._supported_formats[format]
image_generator(smiles, image)
[docs] def register(self, format, generator):
self._supported_formats[format] = generator
[docs]class SchrodImageGenerator(ImageGenerator):
"""
A class to generate 2D structure image file from SMILES using Schrodinger
toolkit.
"""
[docs] def __init__(self):
ImageGenerator.__init__(self)
self.virtual_x = None
self._load_modules()
self._app = self.QtWidgets.QApplication.instance()
if not self._app:
self._app = self.QtWidgets.QApplication([])
self._model = self.canvas2d.ChmRender2DModel()
self._renderer = self.canvas2d.Chm2DRenderer(self._model)
self.register('bmp', self.smiles2image)
self.register('jpg', self.smiles2image)
self.register('jpeg', self.smiles2image)
self.register('png', self.smiles2image)
self.register('ppm', self.smiles2image)
self.register('tiff', self.smiles2image)
self.register('xbm', self.smiles2image)
self.register('xpm', self.smiles2image)
self.register('svg', self.smiles2svg)
self._coor_gen = True
[docs] def setBoundingRect(self, x, y, width, height):
renderer_rect = self.canvas2d.ChmRect2D(x, y, width, height)
self._renderer.setBoundingBox(renderer_rect)
def _load_modules(self):
if sys.platform.startswith('linux') and 'DISPLAY' not in os.environ:
# If there is no DISPLAY variable but there are tests that need X,
# start Xvfb.
self.virtual_x = Xvfb()
try:
self.virtual_x.start()
except OSError:
return
from schrodinger.application.canvas.base import ChmLicenseShared_isValid
from schrodinger.application.canvas.utils import get_license
from schrodinger.infra import canvas2d
from schrodinger.Qt import QtCore
from schrodinger.Qt import QtGui
from schrodinger.Qt import QtSvg
from schrodinger.Qt import QtWidgets
if not ChmLicenseShared_isValid():
self.canvas_license = get_license("LICENSE_SHARED")
self.QtGui = QtGui
self.QtCore = QtCore
self.QtSvg = QtSvg
self.QtWidgets = QtWidgets
self.canvas2d = canvas2d
[docs] def smiles2mol(self, smiles):
return self.canvas2d.ChmMol.fromSMILES(smiles)
[docs] def mol2svg(self, mol, svg_fname):
#FIXME: add an extra argument to mol2svg
#pic = self._renderer.getQPicture(mol, False)
pic = self._renderer.getQPicture(mol, self._coor_gen)
rect = pic.boundingRect()
svg_gen = self.QtSvg.QSvgGenerator()
svg_gen.setFileName(svg_fname)
svg_gen.setSize(self.QtCore.QSize(rect.width(), rect.height()))
#image.fill(QtCore.Qt.white)
qp = self.QtGui.QPainter(svg_gen)
qp.drawPicture(-pic.boundingRect().left(), -pic.boundingRect().top(),
pic)
qp.end()
[docs] def smiles2svg(self, smiles, svg_fname):
mol = self.smiles2mol(smiles)
self.mol2svg(mol, svg_fname)
[docs] def mol2image(
self,
mol,
image_fname,
hl_atoms=[], # noqa: M511
hl_bonds=[]): # noqa: M511
if hl_atoms or hl_bonds:
pic = self._renderer.getQPicture(mol, hl_atoms, hl_bonds,
self.QtGui.QColor(255, 0, 0),
False)
else:
pic = self._renderer.getQPicture(mol, self._coor_gen)
rect = pic.boundingRect()
qimage = self.QtGui.QImage(rect.width(), rect.height(),
self.QtGui.QImage.Format_ARGB32)
qimage.fill(self.QtCore.Qt.white)
qp = self.QtGui.QPainter(qimage)
qp.drawPicture(-pic.boundingRect().left(), -pic.boundingRect().top(),
pic)
qp.end()
qimage.save(image_fname)
[docs] def smiles2image(self, smiles, image_fname):
mol = self.smiles2mol(smiles)
self.mol2image(mol, image_fname)
[docs] def setCoorGenMode(self, mode):
self._coor_gen = mode