implementing app

app.py

@@ -24,6 +24,15 @@ from rdkit import Chem
 from rdkit.Chem import AllChem
 from rdkit import DataStructs
 from PIL import Image
+import matplotlib.pyplot as plt
+
+def plot_bbox(bbox_XYXY, label):
+    xmin, ymin, xmax, ymax =bbox_XYXY
+    plt.plot(
+        [xmin, xmin, xmax, xmax, xmin],
+        [ymin, ymax, ymax, ymin, ymin],
+        color=colors[label],
+        label=str(label))
 
 model_cls = RCNN
 experiment_path_atoms="./models/atoms_model/"
@@ -54,11 +63,10 @@ if image_file is not None:
    trainer = pl.Trainer(logger=False)
    atom_preds = trainer.predict(model_atom, dataset.test_dataloader())
    st.write(atom_preds)
-   #plt.imshow(im, cmap="gray")
-   #     for bbox, label in zip(bboxes_XYXY, labels):
-   #         plot_bbox(bbox, label)
-   #     plt.axis('off')
-   #     plt.savefig("example_image.png",bbox_inches='tight', pad_inches=0)
-   #     plt.show()
+   plt.imshow(image, cmap="gray")
+   for bbox, label in zip(atom_preds[0]['boxes'], atom_preds[0]['pred']):
+       plot_bbox(bbox, label)
+  
+   col2.image(plt.show(), use_column_width=True)
 #x = st.slider('Select a value')
 #st.write(x, 'squared is', x * x)