{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# About This Notebook\n\nThis notebook briefly explores allowed external InChIs and how to use `rdkit` to generate additional training images. It is similar to [this notebook](https://www.kaggle.com/stainsby/improved-synthetic-data-for-bms-competition-v3) except that I try to generate images that more closely resemble those given to us, and with approved InChIs.\n\nWays to use this new data:\n* Pretrain a model with it and finetune on given training data\n* Concatenate with given training data and train model\n* Upsample existing data","metadata":{}},{"cell_type":"markdown","source":"# Library","metadata":{}},{"cell_type":"code","source":"!conda install -y -c rdkit rdkit","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport re\nfrom tqdm.auto import tqdm\ntqdm.pandas()\n\nimport torch\nfrom torch.utils.data import Dataset\nimport cv2\n\nfrom albumentations import (\n    Compose, OneOf, Normalize, Resize, HorizontalFlip, VerticalFlip, Rotate, RandomRotate90, CenterCrop\n    )\n\nfrom albumentations.pytorch import ToTensorV2\n\nimport rdkit.Chem as Chem\n\nfrom matplotlib import pyplot as plt","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Helper Functions\n\nTokenizer class and preprocessing is taken from [@yasufuminakama](https://www.kaggle.com/yasufuminakama)'s notebook [here](https://www.kaggle.com/yasufuminakama/inchi-preprocess-2).","metadata":{}},{"cell_type":"code","source":"class Tokenizer(object):\n    \n    def __init__(self):\n        self.stoi = {}\n        self.itos = {}\n\n    def __len__(self):\n        return len(self.stoi)\n    \n    def fit_on_texts(self, texts):\n        vocab = set()\n        for text in texts:\n            vocab.update(text.split(' '))\n        vocab = sorted(vocab)\n        vocab.append('<sos>')\n        vocab.append('<eos>')\n        vocab.append('<pad>')\n        for i, s in enumerate(vocab):\n            self.stoi[s] = i\n        self.itos = {item[1]: item[0] for item in self.stoi.items()}\n        \n    def text_to_sequence(self, text):\n        sequence = []\n        sequence.append(self.stoi['<sos>'])\n        for s in text.split(' '):\n            sequence.append(self.stoi[s])\n        sequence.append(self.stoi['<eos>'])\n        return sequence\n    \n    def texts_to_sequences(self, texts):\n        sequences = []\n        for text in texts:\n            sequence = self.text_to_sequence(text)\n            sequences.append(sequence)\n        return sequences\n\n    def sequence_to_text(self, sequence):\n        return ''.join(list(map(lambda i: self.itos[i], sequence)))\n    \n    def sequences_to_texts(self, sequences):\n        texts = []\n        for sequence in sequences:\n            text = self.sequence_to_text(sequence)\n            texts.append(text)\n        return texts\n    \n    def predict_caption(self, sequence):\n        caption = ''\n        for i in sequence:\n            if i == self.stoi['<eos>'] or i == self.stoi['<pad>']:\n                break\n            caption += self.itos[i]\n        return caption\n    \n    def predict_captions(self, sequences):\n        captions = []\n        for sequence in sequences:\n            caption = self.predict_caption(sequence)\n            captions.append(caption)\n        return captions\n    \ndef split_form(form):\n    string = ''\n    for i in re.findall(r\"[A-Z][^A-Z]*\", form):\n        elem = re.match(r\"\\D+\", i).group()\n        num = i.replace(elem, \"\")\n        if num == \"\":\n            string += f\"{elem} \"\n        else:\n            string += f\"{elem} {str(num)} \"\n    return string.rstrip(' ')\n\ndef split_form2(form):\n    string = ''\n    for i in re.findall(r\"[a-z][^a-z]*\", form):\n        elem = i[0]\n        num = i.replace(elem, \"\").replace('/', \"\")\n        num_string = ''\n        for j in re.findall(r\"[0-9]+[^0-9]*\", num):\n            num_list = list(re.findall(r'\\d+', j))\n            assert len(num_list) == 1, f\"len(num_list) != 1\"\n            _num = num_list[0]\n            if j == _num:\n                num_string += f\"{_num} \"\n            else:\n                extra = j.replace(_num, \"\")\n                num_string += f\"{_num} {' '.join(list(extra))} \"\n        string += f\"/{elem} {num_string}\"\n    return string.rstrip(' ')\n\ndef get_atom_counts(dataframe):\n    # https://www.kaggle.com/ttahara/bms-mt-chemical-formula-regression-training\n    TARGETS = [\n    'B', 'Br', 'C', 'Cl',\n    'F', 'H', 'I', 'N',\n    'O', 'P', 'S', 'Si']\n    elem_regex = re.compile(r\"[A-Z][a-z]?[0-9]*\")\n    atom_regex = re.compile(r\"[A-Z][a-z]?\")\n    dgts_regex = re.compile(r\"[0-9]*\")\n    \n    atom_dict_list = []\n    for fml in tqdm(dataframe[\"InChI_1\"].values):\n        atom_dict = dict()\n        for elem in elem_regex.findall(fml):\n            atom = dgts_regex.sub(\"\", elem)\n            dgts = atom_regex.sub(\"\", elem)\n            atom_cnt = int(dgts) if len(dgts) > 0 else 1\n            atom_dict[atom] = atom_cnt\n        atom_dict_list.append(atom_dict)\n\n    atom_df = pd.DataFrame(\n        atom_dict_list).fillna(0).astype(int)\n    atom_df = atom_df.sort_index(axis=\"columns\")\n    for atom in TARGETS:\n        dataframe[atom] = atom_df[atom]\n    return dataframe","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_train_file_path(image_id):\n    return \"../input/bms-molecular-translation/train/{}/{}/{}/{}.png\".format(\n        image_id[0], image_id[1], image_id[2], image_id \n    )","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data Loading","metadata":{}},{"cell_type":"code","source":"OUTPUT_DIR = './'","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train = pd.read_pickle('../input/inchi-preprocess-2/train2.pkl')\ntrain['file_path'] = train['image_id'].apply(get_train_file_path)\nallowed_inchi = pd.read_csv('../input/bms-molecular-translation/extra_approved_InChIs.csv')\nprint(allowed_inchi.shape)\nallowed_inchi.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"allowed_inchi['InChI_1'] = allowed_inchi['InChI'].progress_apply(lambda x: x.split('/')[1])\nallowed_inchi['InChI_text'] = allowed_inchi['InChI_1'].progress_apply(split_form) + ' ' + \\\n                        allowed_inchi['InChI'].apply(lambda x: '/'.join(x.split('/')[2:])).progress_apply(split_form2).values","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train, allowed_inchi = get_atom_counts(train), get_atom_counts(allowed_inchi)\nallowed_inchi.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# ====================================================\n# create tokenizer\n# ====================================================\ntokenizer = Tokenizer()\ntokenizer.fit_on_texts(allowed_inchi['InChI_text'].values)\ntorch.save(tokenizer, 'tokenizer2.pth')\nprint('Saved tokenizer')\n# ====================================================\n# preprocess allowed_inchi.csv\n# ====================================================\nlengths = []\nfor text in tqdm(allowed_inchi['InChI_text'].values, total=len(allowed_inchi)):\n    seq = tokenizer.text_to_sequence(text)\n    length = len(seq) - 2\n    lengths.append(length)\nallowed_inchi['InChI_length'] = lengths\nallowed_inchi.to_pickle('allowed_inchi_processed.pkl')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Brief EDA","metadata":{}},{"cell_type":"code","source":"# compare length distribution to given training sequences\nprint(allowed_inchi['InChI_length'].min(), allowed_inchi['InChI_length'].max())\nplt.figure(figsize=(15, 5))\nplt.subplot(1, 2, 1)\nplt.title('External')\nplt.hist(allowed_inchi['InChI_length'].values, bins=20)\nplt.subplot(1, 2, 2)\nplt.title('Given')\nplt.hist(train['InChI_length'].values, bins=20);","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# ensure no overlap between inchis in train.csv\nallowed_inchi[allowed_inchi['InChI_text'].isin(train['InChI_text']).values]","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# ensure no duplicate inchis\nallowed_inchi['InChI'].nunique() == len(allowed_inchi)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# explore atom count distribution\natoms = [\n    'B', 'Br', 'C', 'Cl',\n    'F', 'H', 'I', 'N',\n    'O', 'P', 'S', 'Si']\nfig, ax = plt.subplots(len(atoms), 2, figsize=(15, 20))\nfor i, atom in enumerate(atoms):\n    ax[i, 0].hist(allowed_inchi[atom].values, bins=20, density=True)\n    ax[i, 0].set_title(f'External - {atom}')\n    ax[i, 1].hist(train[atom].values, bins=20, density=True)\n    ax[i, 1].set_title(f'Given - {atom}')\n    fig.tight_layout()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Generate Additional Training Images","metadata":{}},{"cell_type":"code","source":"def sp_noise(image):\n    #https://gist.github.com/lucaswiman/1e877a164a69f78694f845eab45c381a\n    output = image.copy()\n    if len(image.shape) == 2:\n        black = 0\n        white = 255            \n    else:\n        colorspace = image.shape[2]\n        if colorspace == 3:  # RGB\n            black = np.array([0, 0, 0], dtype='uint8')\n            white = np.array([255, 255, 255], dtype='uint8')\n        else:  # RGBA\n            black = np.array([0, 0, 0, 255], dtype='uint8')\n            white = np.array([255, 255, 255, 255], dtype='uint8')\n    probs = np.random.random(image.shape[:2])\n    image[probs < .00015] = black\n    image[probs > .85] = white\n    return image\n\ndef noisy_inchi(inchi, inchi_path, add_noise=True, crop_and_pad=True):\n    mol = Chem.MolFromInchi(inchi)\n    d = Chem.Draw.rdMolDraw2D.MolDraw2DCairo(300, 300)\n    # https://www.kaggle.com/stainsby/improved-synthetic-data-for-bms-competition-v3\n    Chem.rdDepictor.SetPreferCoordGen(True)\n    d.drawOptions().maxFontSize=14\n    d.drawOptions().multipleBondOffset=np.random.uniform(0.05, 0.2)\n    d.drawOptions().useBWAtomPalette()\n    d.drawOptions().bondLineWidth=1\n    d.drawOptions().additionalAtomLabelPadding=np.random.uniform(0, .2)\n    d.DrawMolecule(mol)\n    d.FinishDrawing()\n    d.WriteDrawingText(inchi_path)  \n    if crop_and_pad:\n        img = cv2.imread(inchi_path, cv2.IMREAD_GRAYSCALE)\n        crop_rows = img[~np.all(img==255, axis=1), :]\n        img = crop_rows[:, ~np.all(crop_rows==255, axis=0)]\n        img = cv2.copyMakeBorder(img, 30, 30, 30, 30, cv2.BORDER_CONSTANT, value=255)\n        img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)\n    else:\n        img = cv2.imread(inchi_path)\n    if add_noise:\n        img = sp_noise(img)\n        cv2.imwrite(inchi_path, img)\n    return img","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# ====================================================\n# Dataset\n# ====================================================\nclass ExternalDataset(Dataset):\n    def __init__(self, df, tokenizer, transform=None):\n        super().__init__()\n        self.df = df\n        self.tokenizer = tokenizer\n        self.labels = df['InChI_text'].values\n        self.inchis = df['InChI'].values\n        self.transform = transform\n    \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, idx):\n        inchi = self.inchis[idx]\n        image = noisy_inchi(inchi, inchi_path=f'{OUTPUT_DIR}{idx}.png')\n        if self.transform:\n            augmented = self.transform(image=image)\n            image = augmented['image']\n        label = self.labels[idx]\n        label = self.tokenizer.text_to_sequence(label)\n        label_length = len(label)\n        label_length = torch.LongTensor([label_length])\n        return image, torch.LongTensor(label), label_length\n    \nclass TrainDataset(Dataset):\n    def __init__(self, df, tokenizer, transform=None):\n        super().__init__()\n        self.df = df\n        self.tokenizer = tokenizer\n        self.file_paths = df['file_path'].values\n        self.labels = df['InChI_text'].values\n        self.transform = transform\n    \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, idx):\n        file_path = self.file_paths[idx]\n        image = cv2.imread(file_path)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        if self.transform:\n            augmented = self.transform(image=image)\n            image = augmented['image']\n        label = self.labels[idx]\n        label = self.tokenizer.text_to_sequence(label)\n        label_length = len(label)\n        label_length = torch.LongTensor([label_length])\n        return image, torch.LongTensor(label), label_length\n    \n# ====================================================\n# Augmentations\n# ====================================================\ndef get_transforms(*, data):\n    if data == 'train':\n        return Compose([\n            Resize(224, 224),\n            OneOf([\n                 VerticalFlip(),\n                 RandomRotate90(),\n                  ], p=.5),\n            Normalize(\n                mean=[0.485, 0.456, 0.406],\n                std=[0.229, 0.224, 0.225],\n            ),\n            ToTensorV2(),\n        ])\n    \n    elif data == 'valid':\n        return Compose([\n            Resize(224, 224),\n            Normalize(\n                mean=[0.485, 0.456, 0.406],\n                std=[0.229, 0.224, 0.225],\n            ),\n            ToTensorV2(),\n        ])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"start = np.random.randint(100000)\next_dataset = ExternalDataset(allowed_inchi, tokenizer, transform=get_transforms(data='valid'))\n\nfor i in range(start, start+5):\n    plt.figure(figsize=(5, 5))\n    image, label, label_length = ext_dataset[i]\n    text = tokenizer.sequence_to_text(label.numpy())\n    plt.imshow(image.transpose(0, 1).transpose(1, 2))\n    plt.show() ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Compare to Original Images","metadata":{}},{"cell_type":"code","source":"#start = np.random.randint(100000)\nstart=0\next_dataset = ExternalDataset(train, tokenizer, transform=get_transforms(data='valid'))\norig_dataset = TrainDataset(train, tokenizer, transform=get_transforms(data='valid'))\n\nfor i in range(start, start+8):\n    plt.figure(figsize=(10, 10))\n    image, label, label_length = ext_dataset[i]\n    orig_image, _, _ = orig_dataset[i]\n    text = tokenizer.sequence_to_text(label.numpy())\n    plt.subplot(1, 2, 1)\n    plt.title('Generated')\n    plt.imshow(image.transpose(0, 1).transpose(1, 2))\n    plt.subplot(1, 2, 2)\n    plt.title('Given')\n    plt.imshow(orig_image.transpose(0, 1).transpose(1, 2))\n    plt.show() ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Not too bad. You can increase the noise and change the drawing parameters in `rdkit` for further improvements.","metadata":{}}]}