.

app.py

@@ -17,6 +17,7 @@ from dcmnet.modules import MessagePassingModel
 
 from rdkit import Chem
 from rdkit.Chem import AllChem
+from rdkit.Chem import Draw
 
 def get_grid_points(coordinates):
     """
@@ -58,7 +59,7 @@ for i, atom in enumerate(smiles_mol.GetAtoms()):
     # For each atom, set the property "molAtomMapNumber" to a custom number, let's say, the index of the atom in the molecule
     atom.SetProp("atomNote", str(atom.GetIdx()))
 
-smiles_image = Chem.Draw.MolToImage(smiles_mol)
+smiles_image = Draw.MolToImage(smiles_mol)
 
 # display molecule
 st.image(smiles_image)