moldenhof commited on
Commit
648d955
1 Parent(s): 692229b

implementing app

Browse files
Files changed (1) hide show
  1. app.py +14 -6
app.py CHANGED
@@ -24,6 +24,15 @@ from rdkit import Chem
24
  from rdkit.Chem import AllChem
25
  from rdkit import DataStructs
26
  from PIL import Image
 
 
 
 
 
 
 
 
 
27
 
28
  model_cls = RCNN
29
  experiment_path_atoms="./models/atoms_model/"
@@ -54,11 +63,10 @@ if image_file is not None:
54
  trainer = pl.Trainer(logger=False)
55
  atom_preds = trainer.predict(model_atom, dataset.test_dataloader())
56
  st.write(atom_preds)
57
- #plt.imshow(im, cmap="gray")
58
- # for bbox, label in zip(bboxes_XYXY, labels):
59
- # plot_bbox(bbox, label)
60
- # plt.axis('off')
61
- # plt.savefig("example_image.png",bbox_inches='tight', pad_inches=0)
62
- # plt.show()
63
  #x = st.slider('Select a value')
64
  #st.write(x, 'squared is', x * x)
 
24
  from rdkit.Chem import AllChem
25
  from rdkit import DataStructs
26
  from PIL import Image
27
+ import matplotlib.pyplot as plt
28
+
29
+ def plot_bbox(bbox_XYXY, label):
30
+ xmin, ymin, xmax, ymax =bbox_XYXY
31
+ plt.plot(
32
+ [xmin, xmin, xmax, xmax, xmin],
33
+ [ymin, ymax, ymax, ymin, ymin],
34
+ color=colors[label],
35
+ label=str(label))
36
 
37
  model_cls = RCNN
38
  experiment_path_atoms="./models/atoms_model/"
 
63
  trainer = pl.Trainer(logger=False)
64
  atom_preds = trainer.predict(model_atom, dataset.test_dataloader())
65
  st.write(atom_preds)
66
+ plt.imshow(image, cmap="gray")
67
+ for bbox, label in zip(atom_preds[0]['boxes'], atom_preds[0]['pred']):
68
+ plot_bbox(bbox, label)
69
+
70
+ col2.image(plt.show(), use_column_width=True)
 
71
  #x = st.slider('Select a value')
72
  #st.write(x, 'squared is', x * x)