improving layout

app.py

@@ -44,56 +44,35 @@ def plot_bbox(bbox_XYXY, label):
 def atomlenz(modelfile):
     model_cls = RCNN
     experiment_path_atoms="./models/atoms_model/"
-    #dir_list = os.listdir(experiment_path_atoms)
-    #dir_list = [os.path.join(experiment_path_atoms,f) for f in dir_list]
-    #dir_list.sort(key=os.path.getctime, reverse=True)
-    #checkpoint_file_atoms = [f for f in dir_list if "ckpt" in f][0]
     checkpoint_file_atoms=os.path.join(experiment_path_atoms,modelfile)
     model_atom = model_cls.load_from_checkpoint(checkpoint_file_atoms)
     model_atom.model.roi_heads.score_thresh = 0.65
     experiment_path_bonds = "./models/bonds_model/"
-    #dir_list = os.listdir(experiment_path_bonds)
-    #dir_list = [os.path.join(experiment_path_bonds,f) for f in dir_list]
-    #dir_list.sort(key=os.path.getctime, reverse=True)
-    #checkpoint_file_bonds = [f for f in dir_list if "ckpt" in f][0]
     checkpoint_file_bonds=os.path.join(experiment_path_bonds,modelfile)
     model_bond = model_cls.load_from_checkpoint(checkpoint_file_bonds)
     model_bond.model.roi_heads.score_thresh = 0.65
     experiment_path_stereo = "./models/stereos_model/"
-    #dir_list = os.listdir(experiment_path_stereo)
-    #dir_list = [os.path.join(experiment_path_stereo,f) for f in dir_list]
-    #dir_list.sort(key=os.path.getctime, reverse=True)
-    #checkpoint_file_stereo = [f for f in dir_list if "ckpt" in f][0]
     checkpoint_file_stereo=os.path.join(experiment_path_stereo,modelfile)
     model_stereo = model_cls.load_from_checkpoint(checkpoint_file_stereo)
     model_stereo.model.roi_heads.score_thresh = 0.65
     experiment_path_charges = "./models/charges_model/"
-    #dir_list = os.listdir(experiment_path_charges)
-    #dir_list = [os.path.join(experiment_path_charges,f) for f in dir_list]
-    #dir_list.sort(key=os.path.getctime, reverse=True)
-    #checkpoint_file_charges = [f for f in dir_list if "ckpt" in f][0]
     checkpoint_file_charges=os.path.join(experiment_path_charges,modelfile)
     model_charge = model_cls.load_from_checkpoint(checkpoint_file_charges)
     model_charge.model.roi_heads.score_thresh = 0.65
 
     data_cls = Objects_Smiles
     dataset = data_cls(data_path="./uploads/", batch_size=1)
-#    dataset.prepare_data()
 
     image_file = st.file_uploader("Upload a chemical structure candidate image",type=['png'])
-#st.write('filename is', file_name)
     if image_file is not None:
-   #col1, col2 = st.columns(2)
 
        image = Image.open(image_file)
-   #col1.image(image, use_column_width=True)
        st.image(image, use_column_width=True)
        col1, col2 = st.columns(2)
        if not os.path.exists("uploads/images"):
            os.makedirs("uploads/images")
        with open(os.path.join("uploads/images/","0.png"),"wb") as f:
             f.write(image_file.getbuffer())
-   #st.success("Saved File")
        dataset.prepare_data()
        trainer = pl.Trainer(logger=False)
        st.toast('Predicting atoms,bonds,charges,..., please wait')
@@ -102,26 +81,21 @@ def atomlenz(modelfile):
        stereo_preds = trainer.predict(model_stereo, dataset.test_dataloader())
        charges_preds = trainer.predict(model_charge, dataset.test_dataloader())
        st.toast('Done')
-   #st.write(atom_preds)
        plt.imshow(image, cmap="gray")
        for bbox, label in zip(atom_preds[0]['boxes'][0], atom_preds[0]['preds'][0]):
-      # st.write(bbox)
-      # st.write(label)
            plot_bbox(bbox, label)
        plt.axis('off')
        plt.savefig("example_image.png",bbox_inches='tight', pad_inches=0)
        image_vis = Image.open("example_image.png")
-       col1.image(image_vis, use_column_width=True)
+       col1.image(image_vis, caption=f"Atom entities", use_column_width=True)
        plt.clf()
        plt.imshow(image, cmap="gray")
        for bbox, label in zip(bond_preds[0]['boxes'][0], bond_preds[0]['preds'][0]):
-      # st.write(bbox)
-      # st.write(label)
            plot_bbox(bbox, label)
        plt.axis('off')
        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)
+       col2.image(image_vis, caption=f"Bond entities", use_column_width=True)
        mol_graphs = []
        count_bonds_preds = np.zeros(4)
        count_atoms_preds = np.zeros(15)
@@ -142,11 +116,7 @@ def atomlenz(modelfile):
             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))
@@ -162,10 +132,8 @@ def atomlenz(modelfile):
                      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]
@@ -174,14 +142,12 @@ def atomlenz(modelfile):
                      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
             stereo_bonds = np.where(mol_graph>4, True, False)
@@ -198,7 +164,6 @@ def atomlenz(modelfile):
 
             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
@@ -212,7 +177,6 @@ def atomlenz(modelfile):
               if len(problems) > 0:
                  mol = solve_mol_problems(mol,problems)
                  problematic = 1
-           #import ipdb; ipdb.set_trace()
               try:
                 Chem.SanitizeMol(mol)
               except:
@@ -232,7 +196,6 @@ def atomlenz(modelfile):
               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)