implementing app

app.py

@@ -43,6 +43,28 @@ dir_list.sort(key=os.path.getctime, reverse=True)
 checkpoint_file_atoms = [f for f in dir_list if "ckpt" in f][0]
 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]
+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]
+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]
+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()
@@ -65,6 +87,9 @@ if image_file is not None:
    dataset.prepare_data()
    trainer = pl.Trainer(logger=False)
    atom_preds = trainer.predict(model_atom, dataset.test_dataloader())
+   bond_preds = trainer.predict(model_bond, dataset.test_dataloader())
+   stereo_preds = trainer.predict(model_stereo, dataset.test_dataloader())
+   charge_preds = trainer.predict(model_charge, dataset.test_dataloader())
    #st.write(atom_preds)
    plt.imshow(image, cmap="gray")
    for bbox, label in zip(atom_preds[0]['boxes'][0], atom_preds[0]['preds'][0]):
@@ -75,5 +100,15 @@ if image_file is not None:
    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)
+   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)
 #x = st.slider('Select a value')
 #st.write(x, 'squared is', x * x)