implementing predict smiles

app.py

@@ -5,7 +5,7 @@ import torch
 import numpy as np
 #import matplotlib.pyplot as plt
 #import pathlib
-#from AtomLenz import *
+from AtomLenz import *
 #from utils_graph import *
 from Object_Smiles import Objects_Smiles 
 
@@ -112,5 +112,126 @@ if image_file is not None:
    plt.savefig("example_image.png",bbox_inches='tight', pad_inches=0)
    image_vis = Image.open("example_image.png")
    col2.image(image_vis, use_column_width=True)
+   for image_idx, bonds in enumerate(bond_preds):
+        count_bonds_preds = np.zeros(8)
+        count_atoms_preds = np.zeros(18)
+        atom_boxes = atom_preds[image_idx]['boxes'][0]
+        atom_labels = atom_preds[image_idx]['preds'][0]
+        atom_scores = atom_preds[image_idx]['scores'][0]
+        charge_boxes = charges_preds[image_idx]['boxes'][0]
+        charge_labels = charges_preds[image_idx]['preds'][0]
+        charge_mask=torch.where(charge_labels>1)
+        filtered_ch_labels=charge_labels[charge_mask]
+        filtered_ch_boxes=charge_boxes[charge_mask]
+        #import ipdb; ipdb.set_trace()
+        filtered_bboxes, filtered_labels = iou_filter_bboxes(atom_boxes, atom_labels, atom_scores)
+        #for atom_label in filtered_labels:
+        #    count_atoms_preds[atom_label] += 1
+        #import ipdb; ipdb.set_trace()
+        mol_graph = np.zeros((len(filtered_bboxes),len(filtered_bboxes)))
+        stereo_atoms = np.zeros(len(filtered_bboxes))
+        charge_atoms = np.ones(len(filtered_bboxes))
+        for index,box_atom in enumerate(filtered_bboxes):
+            for box_charge,label_charge in zip(filtered_ch_boxes,filtered_ch_labels):
+                if bb_box_intersects(box_atom,box_charge) == 1:
+                    charge_atoms[index]=label_charge
+            
+        for bond_idx, bond_box in enumerate(bonds['boxes'][0]):
+            label_bond = bonds['preds'][0][bond_idx]
+            if label_bond > 1:
+              try:
+                 count_bonds_preds[label_bond] += 1
+              except:
+                 count_bonds_preds=count_bonds_preds 
+               #import ipdb; ipdb.set_trace()
+              result = []
+              limit = 0
+            #TODO: values of 50 and 5 should be made dependent of mean size of atom_boxes
+              while result.count(1) < 2 and limit < 80:
+                 result=[]
+                 bigger_bond_box = [bond_box[0]-limit,bond_box[1]-limit,bond_box[2]+limit,bond_box[3]+limit]
+                 for atom_box in filtered_bboxes:
+                     result.append(bb_box_intersects(atom_box,bigger_bond_box))
+                 limit+=5
+              indices = [i for i, x in enumerate(result) if x == 1]
+              if len(indices) == 2:
+               #import ipdb; ipdb.set_trace()
+                 mol_graph[indices[0],indices[1]]=label_bond
+                 mol_graph[indices[1],indices[0]]=label_bond
+              if len(indices) > 2:
+                #we have more then two canidate atoms for one bond, we filter ...
+                  cand_bboxes = filtered_bboxes[indices,:]
+                  cand_indices = dist_filter_bboxes(cand_bboxes)
+                #import ipdb; ipdb.set_trace()
+                  mol_graph[indices[cand_indices[0]],indices[cand_indices[1]]]=label_bond
+                  mol_graph[indices[cand_indices[1]],indices[cand_indices[0]]]=label_bond
+                  #print("more than 2 indices")
+              #if len(indices) < 2:
+              #    print("less than 2 indices")
+                #import ipdb; ipdb.set_trace()
+ #           else:
+ #             result=[]
+ #             for atom_box in filtered_bboxes:
+ #                 result.append(bb_box_intersects(atom_box,bond_box))
+ #             indices = [i for i, x in enumerate(result) if x == 1]
+ #             if len(indices) == 1:
+ #                stereo_atoms[indices[0]]=label_bond
+        stereo_bonds = np.where(mol_graph>4, True, False)
+        if np.any(stereo_bonds):
+           stereo_boxes = stereo_preds[image_idx]['boxes'][0]
+           stereo_labels= stereo_preds[image_idx]['preds'][0]
+           for stereo_box in stereo_boxes:
+               result=[]
+               for atom_box in filtered_bboxes:
+                   result.append(bb_box_intersects(atom_box,stereo_box))
+               indices = [i for i, x in enumerate(result) if x == 1]
+               if len(indices) == 1:
+                   stereo_atoms[indices[0]]=1
+               
+        molecule = dict()
+        molecule['graph'] = mol_graph
+        #molecule['atom_labels'] = atom_preds[image_idx]['preds'][0]
+        molecule['atom_labels'] = filtered_labels
+        molecule['atom_boxes'] = filtered_bboxes
+        molecule['stereo_atoms'] = stereo_atoms
+        molecule['charge_atoms'] = charge_atoms
+        mol_graphs.append(molecule)
+        base_path = pathlib.Path(args.data_path)
+        image_dir = base_path.joinpath("images")
+        smiles_dir = base_path.joinpath("smiles")
+        impath = image_dir.joinpath(f"{image_idx}.png")
+        smilespath = smiles_dir.joinpath(f"{image_idx}.txt")
+        save_mol_to_file(molecule,'molfile')
+        mol =  Chem.MolFromMolFile('molfile',sanitize=False)
+        problematic = 0
+        try:
+          problems = Chem.DetectChemistryProblems(mol)
+          if len(problems) > 0:
+             mol = solve_mol_problems(mol,problems)
+             problematic = 1
+           #import ipdb; ipdb.set_trace()
+          try:
+            Chem.SanitizeMol(mol) 
+          except:
+            problems = Chem.DetectChemistryProblems(mol)
+            if len(problems) > 0:
+              mol = solve_mol_problems(mol,problems)
+            try:
+              Chem.SanitizeMol(mol)
+            except:
+              pass
+        except:
+          problematic = 1
+        try:
+          pred_smiles = Chem.MolToSmiles(mol)
+        except:
+          pred_smiles = ""
+          problematic = 1
+        predictions+=1
+        predictions_list.append([image_idx,pred_smiles,problematic])
+                #import ipdb; ipdb.set_trace()
+   file_preds = open('preds_atomlenz','w')
+   for pred in predictions_list:
+        print(pred)
 #x = st.slider('Select a value')
 #st.write(x, 'squared is', x * x)