{"cells":[{"metadata":{},"cell_type":"markdown","source":"## Purpose\n\nThis code is designed to generate synthetic data with more random characterics than is provided by the BMS dataset. This could be used to aid the robustness of algorithms created to tackle the BMS task and other related tasks.\n\nFor the background disussion, refer to: https://www.kaggle.com/c/bms-molecular-translation/discussion/225590 .\n\nUpdates for the 2nd edition:\n  - resulting outpus are closer to the original dataset\n  - support InChI inputs (not just SMILES)\n  - font scaling now actually works (in the previous version, additional random diagram scaling made it look as if font scaling was working)\n  - replaced hacky SVG manipulation code for dropping bonds and changing fonts with code using lxml\n  - outputs also include the unaugmented image, given in two layers: one for the atoms and one for the bonds\n  - result are shown actual size and compared with image from the original dataset.\n \nUpdates for 3rd edition:\n  - a slight tweak to stretch the value range of output values of the (unaugmented) atom and bond images to fit exactly in the range 0 to 1"},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip -q install cairosvg==2.5.2\n!pip -q install reportlab==3.5.65\n!pip -q install cssutils==2.2.0","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print('Installing RDKit. This can take a while ...')\n# We use a specific version of RDKit with known working font scaling.\n!conda install -q -y -c rdkit rdkit=2020_03_6\nprint('DONE.')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from pathlib import Path\nfrom io import BytesIO\nimport re\nimport copy\nimport logging\n\nimport numpy as np\nimport pandas as pd\n\nimport lxml.etree as et\nimport cssutils\n\nfrom PIL import Image\nimport cairosvg\nfrom skimage.transform import resize\n\nimport IPython\nfrom IPython.display import SVG\nfrom IPython.display import display\nimport ipywidgets as widgets\nfrom matplotlib import pyplot as plt\n\nimport rdkit\nfrom rdkit import Chem\nfrom rdkit.Chem import AllChem, Draw\nfrom rdkit.Chem.Draw import MolDrawOptions","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"PROJECT_DIR = Path('..')\nINPUT_DIR = PROJECT_DIR / 'input' / 'bms-molecular-translation'\nTMP_DIR = PROJECT_DIR / 'tmp'\nTRAIN_DATA_PATH = INPUT_DIR / 'train'\nTRAIN_LABELS_PATH = INPUT_DIR / 'train_labels.csv'\nTMP_DIR.mkdir(exist_ok=True)\n\ncssutils.log.setLevel(logging.CRITICAL)\n\nnp.set_printoptions(edgeitems=30, linewidth=180)\nprint('RDKit version:', rdkit.__version__)\n# Use a specific version of RDKit with known characteristics so that we can reliably manipulate output SVG.\nassert rdkit.__version__ == '2020.03.6'\n\nTRAIN_LABELS = pd.read_csv(TRAIN_LABELS_PATH)\nprint(f'Read {len(TRAIN_LABELS)} training labels.')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Implementation\n\nWe vary the drawing properties, the font, and occassionally drop bond lines."},{"metadata":{"trusted":true},"cell_type":"code","source":"def one_in(n):\n    return np.random.randint(n) == 0 and True or False\n\n\ndef yesno():\n    return one_in(2)\n\n\ndef svg_to_image(svg, convert_to_greyscale=True):\n    svg_str = et.tostring(svg)\n    # TODO: would prefer to convert SVG dirrectly to a numpy array.\n    png = cairosvg.svg2png(bytestring=svg_str)\n    image = np.array(Image.open(BytesIO(png)), dtype=np.float32)\n    # Naive greyscale conversion.\n    if convert_to_greyscale:\n        image = image.mean(axis=-1)\n    return image\n\n\ndef elemstr(elem):\n    return ', '.join([item[0] + ': ' + item[1] for item in elem.items()])\n\n\n# Streches the value range of an image to be exactly 0, 1, unless the image appears to be blank.\ndef stretch_image(img, blank_threshold=1e-2):\n    img_min = img.min()\n    img = img - img_min\n    img_max = img.max()\n    if img_max < blank_threshold:\n        # seems to be blank or close to it\n        return img\n    img_max = img.max()\n    if img_max < 1.0:\n        img = img/img_max\n    return img\n\n\ndef random_molecule_image(inchi, drop_bonds=True, add_noise=True, render_size=1200, margin_fraction=0.2):\n    # Note that the original image is returned as two layers: one for atoms and one for bonds.\n    #mol = Chem.MolFromSmiles(smiles)\n    mol = Chem.inchi.MolFromInchi(inchi)\n    d = Draw.rdMolDraw2D.MolDraw2DSVG(render_size, render_size)\n    options = MolDrawOptions()\n    options.useBWAtomPalette()\n    options.additionalAtomLabelPadding = np.random.uniform(0, 0.3)\n    options.bondLineWidth = int(np.random.uniform(1, 4))\n    options.multipleBondOffset = np.random.uniform(0.05, 0.2)\n    options.rotate = np.random.uniform(0, 360)\n    options.fixedScale = np.random.uniform(0.05, 0.07)\n    options.minFontSize = 20\n    options.maxFontSize = options.minFontSize + int(np.round(np.random.uniform(0, 36)))\n    d.SetFontSize(100)\n    d.SetDrawOptions(options)\n    d.DrawMolecule(mol)\n    d.FinishDrawing()\n    svg_str = d.GetDrawingText()\n    # Do some SVG manipulation\n    svg = et.fromstring(svg_str.encode('iso-8859-1'))\n    atom_elems = svg.xpath(r'//svg:text', namespaces={'svg': 'http://www.w3.org/2000/svg'})\n    bond_elems = svg.xpath(r'//svg:path[starts-with(@class,\"bond-\")]', namespaces={'svg': 'http://www.w3.org/2000/svg'})\n    # Change the font.\n    font_family = np.random.choice([\n        'serif',\n        'sans-serif'\n    ])\n    for elem in atom_elems:\n        style = elem.attrib['style']\n        css = cssutils.parseStyle(style)\n        css.setProperty('font-family', font_family)\n        css_str = css.cssText.replace('\\n', ' ')\n        elem.attrib['style'] = css_str\n    # Create the original image layers.\n    # TODO: separate atom and bond layers\n    bond_svg = copy.deepcopy(svg)\n    # remove atoms from bond_svg\n    for elem in bond_svg.xpath(r'//svg:text', namespaces={'svg': 'http://www.w3.org/2000/svg'}):\n        parent_elem = elem.getparent()\n        if parent_elem is not None:\n            parent_elem.remove(elem)\n    orig_bond_img = svg_to_image(bond_svg)\n    atom_svg = copy.deepcopy(svg)\n    # remove bonds from atom_svg\n    for elem in atom_svg.xpath(r'//svg:path', namespaces={'svg': 'http://www.w3.org/2000/svg'}):\n        parent_elem = elem.getparent()\n        if parent_elem is not None:\n            parent_elem.remove(elem)\n    orig_atom_img = svg_to_image(atom_svg)\n    if drop_bonds:\n        num_bond_elems = len(bond_elems)\n        if one_in(3):\n            while True:\n                # drop a bond\n                # Let's leave at least one bond!\n                if num_bond_elems > 1:\n                    bond_elem_idx = np.random.randint(num_bond_elems)\n                    bond_elem = bond_elems[bond_elem_idx]\n                    bond_parent_elem = bond_elem.getparent()\n                    if bond_parent_elem is not None:\n                        bond_parent_elem.remove(bond_elem)\n                        num_bond_elems -= 1\n                else:\n                    break\n                if not one_in(4):\n                    break\n    img = svg_to_image(svg) > 254\n    img = 1*img  # bool → int\n    # Calculate the margins.\n    black_indices = np.where(img == 0)\n    row_indices, col_indices = black_indices\n    if len(row_indices) >= 2:\n        min_y, max_y = row_indices.min(), row_indices.max() + 1\n    else:\n        min_y, max_y = 0, render_size\n    if len(col_indices) >= 2:\n        min_x, max_x = col_indices.min(), col_indices.max() + 1\n    else:\n        min_x, max_x = 0, render_size\n    margin_size = int(np.random.uniform(0.8*margin_fraction, 1.2*margin_fraction)*max(max_y - min_y, max_x - min_x))\n    min_y, max_y = max(min_y - margin_size, 0), min(max_y + margin_size, render_size)\n    min_x, max_x = max(min_x - margin_size, 0), min(max_x + margin_size, render_size)\n    img = img[min_y:max_y, min_x:max_x]\n    img = img.reshape([img.shape[0], img.shape[1]]).astype(np.float32)\n    orig_bond_img = orig_bond_img[min_y:max_y, min_x:max_x]\n    orig_atom_img = orig_atom_img[min_y:max_y, min_x:max_x]\n    scale = np.random.uniform(0.2, 0.4)\n    sz = (np.array(orig_bond_img.shape[:2], dtype=np.float32)*scale).astype(np.int32)\n    orig_bond_img = resize(orig_bond_img, sz, anti_aliasing=True)\n    orig_atom_img = resize(orig_atom_img, sz, anti_aliasing=True)\n    img = resize(img, sz, anti_aliasing=False)\n    img = img > 0.5\n    if add_noise:\n        # Add \"salt and pepper\" noise.\n        salt_amount = np.random.uniform(0, 0.3)\n        salt = np.random.uniform(0, 1, img.shape) < salt_amount\n        img = np.logical_or(img, salt)\n        pepper_amount = np.random.uniform(0, 0.001)\n        pepper = np.random.uniform(0, 1, img.shape) < pepper_amount\n        img = np.logical_or(1 - img, pepper)\n    \n    img = img.astype(np.uint8)  # boolean -> uint8\n    orig_bond_img = 1 - orig_bond_img/255\n    orig_atom_img = 1 - orig_atom_img/255\n    # Stretch the range of the atom and bond images so tha tthe min is 0 and the max. is 1\n    orig_bond_img = stretch_image(orig_bond_img)\n    orig_atom_img = stretch_image(orig_atom_img)\n    return img, orig_bond_img, orig_atom_img","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Tests"},{"metadata":{"trusted":true},"cell_type":"code","source":"def image_widget(a, greyscale=True):\n    img_bytes = BytesIO()\n    img_pil = Image.fromarray(a)\n    if greyscale:\n        img_pil = img_pil.convert(\"L\")\n    else:\n        img_pil = img_pil.convert(\"RGB\")\n    img_pil.save(img_bytes, format='PNG')\n    return widgets.Image(value=img_bytes.getvalue())\n\n\ndef test_random_molecule_image(n=4, graphics=True):\n    for imol in range(n):\n        #smiles = np.random.choice(some_smiles)\n        mol_index = np.random.randint(len(TRAIN_LABELS))\n        mol_id, inchi = TRAIN_LABELS['image_id'][mol_index], TRAIN_LABELS['InChI'][mol_index]\n        mol_train_img_path = TRAIN_DATA_PATH / mol_id[0] /mol_id[1] / mol_id[2] / (mol_id + '.png')\n        train_img = Image.open(mol_train_img_path)\n        img, orig_bond_img, orig_atom_img = random_molecule_image(inchi)\n        if graphics:\n            print('+-------------------------------------------------------------------------------')\n            print(f'Molecule #{imol + 1}: {mol_id}: {inchi}')\n            print('Training image path:', mol_train_img_path)\n            print('Size:', img.shape)\n            combined_orig_img =  np.clip(np.stack([orig_atom_img, orig_bond_img, np.zeros_like(orig_bond_img)], axis=-1), 0.0, 1.0)\n            combined_orig_img = (255*combined_orig_img).astype(np.uint8)\n            widget1 = image_widget(combined_orig_img, greyscale=False)\n            widget2 = image_widget((255*(1 - img)).astype(np.uint8))\n            sidebyside = widgets.HBox([widget1, widget2])\n            display(sidebyside)\n            print(f'Image from training data:')\n            print('Size:', train_img.size)\n            display(train_img)\n    return\n\n\n%timeit -n1 -r10 test_random_molecule_image(n=1, graphics=False)\ntest_random_molecule_image(n=20, graphics=True)","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}