improving layout
Browse files
app.py
CHANGED
@@ -44,56 +44,35 @@ def plot_bbox(bbox_XYXY, label):
|
|
44 |
def atomlenz(modelfile):
|
45 |
model_cls = RCNN
|
46 |
experiment_path_atoms="./models/atoms_model/"
|
47 |
-
#dir_list = os.listdir(experiment_path_atoms)
|
48 |
-
#dir_list = [os.path.join(experiment_path_atoms,f) for f in dir_list]
|
49 |
-
#dir_list.sort(key=os.path.getctime, reverse=True)
|
50 |
-
#checkpoint_file_atoms = [f for f in dir_list if "ckpt" in f][0]
|
51 |
checkpoint_file_atoms=os.path.join(experiment_path_atoms,modelfile)
|
52 |
model_atom = model_cls.load_from_checkpoint(checkpoint_file_atoms)
|
53 |
model_atom.model.roi_heads.score_thresh = 0.65
|
54 |
experiment_path_bonds = "./models/bonds_model/"
|
55 |
-
#dir_list = os.listdir(experiment_path_bonds)
|
56 |
-
#dir_list = [os.path.join(experiment_path_bonds,f) for f in dir_list]
|
57 |
-
#dir_list.sort(key=os.path.getctime, reverse=True)
|
58 |
-
#checkpoint_file_bonds = [f for f in dir_list if "ckpt" in f][0]
|
59 |
checkpoint_file_bonds=os.path.join(experiment_path_bonds,modelfile)
|
60 |
model_bond = model_cls.load_from_checkpoint(checkpoint_file_bonds)
|
61 |
model_bond.model.roi_heads.score_thresh = 0.65
|
62 |
experiment_path_stereo = "./models/stereos_model/"
|
63 |
-
#dir_list = os.listdir(experiment_path_stereo)
|
64 |
-
#dir_list = [os.path.join(experiment_path_stereo,f) for f in dir_list]
|
65 |
-
#dir_list.sort(key=os.path.getctime, reverse=True)
|
66 |
-
#checkpoint_file_stereo = [f for f in dir_list if "ckpt" in f][0]
|
67 |
checkpoint_file_stereo=os.path.join(experiment_path_stereo,modelfile)
|
68 |
model_stereo = model_cls.load_from_checkpoint(checkpoint_file_stereo)
|
69 |
model_stereo.model.roi_heads.score_thresh = 0.65
|
70 |
experiment_path_charges = "./models/charges_model/"
|
71 |
-
#dir_list = os.listdir(experiment_path_charges)
|
72 |
-
#dir_list = [os.path.join(experiment_path_charges,f) for f in dir_list]
|
73 |
-
#dir_list.sort(key=os.path.getctime, reverse=True)
|
74 |
-
#checkpoint_file_charges = [f for f in dir_list if "ckpt" in f][0]
|
75 |
checkpoint_file_charges=os.path.join(experiment_path_charges,modelfile)
|
76 |
model_charge = model_cls.load_from_checkpoint(checkpoint_file_charges)
|
77 |
model_charge.model.roi_heads.score_thresh = 0.65
|
78 |
|
79 |
data_cls = Objects_Smiles
|
80 |
dataset = data_cls(data_path="./uploads/", batch_size=1)
|
81 |
-
# dataset.prepare_data()
|
82 |
|
83 |
image_file = st.file_uploader("Upload a chemical structure candidate image",type=['png'])
|
84 |
-
#st.write('filename is', file_name)
|
85 |
if image_file is not None:
|
86 |
-
#col1, col2 = st.columns(2)
|
87 |
|
88 |
image = Image.open(image_file)
|
89 |
-
#col1.image(image, use_column_width=True)
|
90 |
st.image(image, use_column_width=True)
|
91 |
col1, col2 = st.columns(2)
|
92 |
if not os.path.exists("uploads/images"):
|
93 |
os.makedirs("uploads/images")
|
94 |
with open(os.path.join("uploads/images/","0.png"),"wb") as f:
|
95 |
f.write(image_file.getbuffer())
|
96 |
-
#st.success("Saved File")
|
97 |
dataset.prepare_data()
|
98 |
trainer = pl.Trainer(logger=False)
|
99 |
st.toast('Predicting atoms,bonds,charges,..., please wait')
|
@@ -102,26 +81,21 @@ def atomlenz(modelfile):
|
|
102 |
stereo_preds = trainer.predict(model_stereo, dataset.test_dataloader())
|
103 |
charges_preds = trainer.predict(model_charge, dataset.test_dataloader())
|
104 |
st.toast('Done')
|
105 |
-
#st.write(atom_preds)
|
106 |
plt.imshow(image, cmap="gray")
|
107 |
for bbox, label in zip(atom_preds[0]['boxes'][0], atom_preds[0]['preds'][0]):
|
108 |
-
# st.write(bbox)
|
109 |
-
# st.write(label)
|
110 |
plot_bbox(bbox, label)
|
111 |
plt.axis('off')
|
112 |
plt.savefig("example_image.png",bbox_inches='tight', pad_inches=0)
|
113 |
image_vis = Image.open("example_image.png")
|
114 |
-
col1.image(image_vis, use_column_width=True)
|
115 |
plt.clf()
|
116 |
plt.imshow(image, cmap="gray")
|
117 |
for bbox, label in zip(bond_preds[0]['boxes'][0], bond_preds[0]['preds'][0]):
|
118 |
-
# st.write(bbox)
|
119 |
-
# st.write(label)
|
120 |
plot_bbox(bbox, label)
|
121 |
plt.axis('off')
|
122 |
plt.savefig("example_image.png",bbox_inches='tight', pad_inches=0)
|
123 |
image_vis = Image.open("example_image.png")
|
124 |
-
col2.image(image_vis, use_column_width=True)
|
125 |
mol_graphs = []
|
126 |
count_bonds_preds = np.zeros(4)
|
127 |
count_atoms_preds = np.zeros(15)
|
@@ -142,11 +116,7 @@ def atomlenz(modelfile):
|
|
142 |
charge_mask=torch.where(charge_labels>1)
|
143 |
filtered_ch_labels=charge_labels[charge_mask]
|
144 |
filtered_ch_boxes=charge_boxes[charge_mask]
|
145 |
-
#import ipdb; ipdb.set_trace()
|
146 |
filtered_bboxes, filtered_labels = iou_filter_bboxes(atom_boxes, atom_labels, atom_scores)
|
147 |
-
#for atom_label in filtered_labels:
|
148 |
-
# count_atoms_preds[atom_label] += 1
|
149 |
-
#import ipdb; ipdb.set_trace()
|
150 |
mol_graph = np.zeros((len(filtered_bboxes),len(filtered_bboxes)))
|
151 |
stereo_atoms = np.zeros(len(filtered_bboxes))
|
152 |
charge_atoms = np.ones(len(filtered_bboxes))
|
@@ -162,10 +132,8 @@ def atomlenz(modelfile):
|
|
162 |
count_bonds_preds[label_bond] += 1
|
163 |
except:
|
164 |
count_bonds_preds=count_bonds_preds
|
165 |
-
#import ipdb; ipdb.set_trace()
|
166 |
result = []
|
167 |
limit = 0
|
168 |
-
#TODO: values of 50 and 5 should be made dependent of mean size of atom_boxes
|
169 |
while result.count(1) < 2 and limit < 80:
|
170 |
result=[]
|
171 |
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):
|
|
174 |
limit+=5
|
175 |
indices = [i for i, x in enumerate(result) if x == 1]
|
176 |
if len(indices) == 2:
|
177 |
-
#import ipdb; ipdb.set_trace()
|
178 |
mol_graph[indices[0],indices[1]]=label_bond
|
179 |
mol_graph[indices[1],indices[0]]=label_bond
|
180 |
if len(indices) > 2:
|
181 |
#we have more then two canidate atoms for one bond, we filter ...
|
182 |
cand_bboxes = filtered_bboxes[indices,:]
|
183 |
cand_indices = dist_filter_bboxes(cand_bboxes)
|
184 |
-
#import ipdb; ipdb.set_trace()
|
185 |
mol_graph[indices[cand_indices[0]],indices[cand_indices[1]]]=label_bond
|
186 |
mol_graph[indices[cand_indices[1]],indices[cand_indices[0]]]=label_bond
|
187 |
stereo_bonds = np.where(mol_graph>4, True, False)
|
@@ -198,7 +164,6 @@ def atomlenz(modelfile):
|
|
198 |
|
199 |
molecule = dict()
|
200 |
molecule['graph'] = mol_graph
|
201 |
-
#molecule['atom_labels'] = atom_preds[image_idx]['preds'][0]
|
202 |
molecule['atom_labels'] = filtered_labels
|
203 |
molecule['atom_boxes'] = filtered_bboxes
|
204 |
molecule['stereo_atoms'] = stereo_atoms
|
@@ -212,7 +177,6 @@ def atomlenz(modelfile):
|
|
212 |
if len(problems) > 0:
|
213 |
mol = solve_mol_problems(mol,problems)
|
214 |
problematic = 1
|
215 |
-
#import ipdb; ipdb.set_trace()
|
216 |
try:
|
217 |
Chem.SanitizeMol(mol)
|
218 |
except:
|
@@ -232,7 +196,6 @@ def atomlenz(modelfile):
|
|
232 |
problematic = 1
|
233 |
predictions+=1
|
234 |
predictions_list.append([image_idx,pred_smiles,problematic])
|
235 |
-
#import ipdb; ipdb.set_trace()
|
236 |
file_preds = open('preds_atomlenz','w')
|
237 |
for pred in predictions_list:
|
238 |
print(pred)
|
|
|
44 |
def atomlenz(modelfile):
|
45 |
model_cls = RCNN
|
46 |
experiment_path_atoms="./models/atoms_model/"
|
|
|
|
|
|
|
|
|
47 |
checkpoint_file_atoms=os.path.join(experiment_path_atoms,modelfile)
|
48 |
model_atom = model_cls.load_from_checkpoint(checkpoint_file_atoms)
|
49 |
model_atom.model.roi_heads.score_thresh = 0.65
|
50 |
experiment_path_bonds = "./models/bonds_model/"
|
|
|
|
|
|
|
|
|
51 |
checkpoint_file_bonds=os.path.join(experiment_path_bonds,modelfile)
|
52 |
model_bond = model_cls.load_from_checkpoint(checkpoint_file_bonds)
|
53 |
model_bond.model.roi_heads.score_thresh = 0.65
|
54 |
experiment_path_stereo = "./models/stereos_model/"
|
|
|
|
|
|
|
|
|
55 |
checkpoint_file_stereo=os.path.join(experiment_path_stereo,modelfile)
|
56 |
model_stereo = model_cls.load_from_checkpoint(checkpoint_file_stereo)
|
57 |
model_stereo.model.roi_heads.score_thresh = 0.65
|
58 |
experiment_path_charges = "./models/charges_model/"
|
|
|
|
|
|
|
|
|
59 |
checkpoint_file_charges=os.path.join(experiment_path_charges,modelfile)
|
60 |
model_charge = model_cls.load_from_checkpoint(checkpoint_file_charges)
|
61 |
model_charge.model.roi_heads.score_thresh = 0.65
|
62 |
|
63 |
data_cls = Objects_Smiles
|
64 |
dataset = data_cls(data_path="./uploads/", batch_size=1)
|
|
|
65 |
|
66 |
image_file = st.file_uploader("Upload a chemical structure candidate image",type=['png'])
|
|
|
67 |
if image_file is not None:
|
|
|
68 |
|
69 |
image = Image.open(image_file)
|
|
|
70 |
st.image(image, use_column_width=True)
|
71 |
col1, col2 = st.columns(2)
|
72 |
if not os.path.exists("uploads/images"):
|
73 |
os.makedirs("uploads/images")
|
74 |
with open(os.path.join("uploads/images/","0.png"),"wb") as f:
|
75 |
f.write(image_file.getbuffer())
|
|
|
76 |
dataset.prepare_data()
|
77 |
trainer = pl.Trainer(logger=False)
|
78 |
st.toast('Predicting atoms,bonds,charges,..., please wait')
|
|
|
81 |
stereo_preds = trainer.predict(model_stereo, dataset.test_dataloader())
|
82 |
charges_preds = trainer.predict(model_charge, dataset.test_dataloader())
|
83 |
st.toast('Done')
|
|
|
84 |
plt.imshow(image, cmap="gray")
|
85 |
for bbox, label in zip(atom_preds[0]['boxes'][0], atom_preds[0]['preds'][0]):
|
|
|
|
|
86 |
plot_bbox(bbox, label)
|
87 |
plt.axis('off')
|
88 |
plt.savefig("example_image.png",bbox_inches='tight', pad_inches=0)
|
89 |
image_vis = Image.open("example_image.png")
|
90 |
+
col1.image(image_vis, caption=f"Atom entities", use_column_width=True)
|
91 |
plt.clf()
|
92 |
plt.imshow(image, cmap="gray")
|
93 |
for bbox, label in zip(bond_preds[0]['boxes'][0], bond_preds[0]['preds'][0]):
|
|
|
|
|
94 |
plot_bbox(bbox, label)
|
95 |
plt.axis('off')
|
96 |
plt.savefig("example_image.png",bbox_inches='tight', pad_inches=0)
|
97 |
image_vis = Image.open("example_image.png")
|
98 |
+
col2.image(image_vis, caption=f"Bond entities", use_column_width=True)
|
99 |
mol_graphs = []
|
100 |
count_bonds_preds = np.zeros(4)
|
101 |
count_atoms_preds = np.zeros(15)
|
|
|
116 |
charge_mask=torch.where(charge_labels>1)
|
117 |
filtered_ch_labels=charge_labels[charge_mask]
|
118 |
filtered_ch_boxes=charge_boxes[charge_mask]
|
|
|
119 |
filtered_bboxes, filtered_labels = iou_filter_bboxes(atom_boxes, atom_labels, atom_scores)
|
|
|
|
|
|
|
120 |
mol_graph = np.zeros((len(filtered_bboxes),len(filtered_bboxes)))
|
121 |
stereo_atoms = np.zeros(len(filtered_bboxes))
|
122 |
charge_atoms = np.ones(len(filtered_bboxes))
|
|
|
132 |
count_bonds_preds[label_bond] += 1
|
133 |
except:
|
134 |
count_bonds_preds=count_bonds_preds
|
|
|
135 |
result = []
|
136 |
limit = 0
|
|
|
137 |
while result.count(1) < 2 and limit < 80:
|
138 |
result=[]
|
139 |
bigger_bond_box = [bond_box[0]-limit,bond_box[1]-limit,bond_box[2]+limit,bond_box[3]+limit]
|
|
|
142 |
limit+=5
|
143 |
indices = [i for i, x in enumerate(result) if x == 1]
|
144 |
if len(indices) == 2:
|
|
|
145 |
mol_graph[indices[0],indices[1]]=label_bond
|
146 |
mol_graph[indices[1],indices[0]]=label_bond
|
147 |
if len(indices) > 2:
|
148 |
#we have more then two canidate atoms for one bond, we filter ...
|
149 |
cand_bboxes = filtered_bboxes[indices,:]
|
150 |
cand_indices = dist_filter_bboxes(cand_bboxes)
|
|
|
151 |
mol_graph[indices[cand_indices[0]],indices[cand_indices[1]]]=label_bond
|
152 |
mol_graph[indices[cand_indices[1]],indices[cand_indices[0]]]=label_bond
|
153 |
stereo_bonds = np.where(mol_graph>4, True, False)
|
|
|
164 |
|
165 |
molecule = dict()
|
166 |
molecule['graph'] = mol_graph
|
|
|
167 |
molecule['atom_labels'] = filtered_labels
|
168 |
molecule['atom_boxes'] = filtered_bboxes
|
169 |
molecule['stereo_atoms'] = stereo_atoms
|
|
|
177 |
if len(problems) > 0:
|
178 |
mol = solve_mol_problems(mol,problems)
|
179 |
problematic = 1
|
|
|
180 |
try:
|
181 |
Chem.SanitizeMol(mol)
|
182 |
except:
|
|
|
196 |
problematic = 1
|
197 |
predictions+=1
|
198 |
predictions_list.append([image_idx,pred_smiles,problematic])
|
|
|
199 |
file_preds = open('preds_atomlenz','w')
|
200 |
for pred in predictions_list:
|
201 |
print(pred)
|