Update app.py

app.py

@@ -20,7 +20,6 @@ idx_to_labels = {0:'other',1:'C',2:'O',3:'N',4:'Cl',5:'Br',6:'S',7:'F',8:'B',
 
 def image_to_numpy(image_path):
 
-    # image = Image.open(image_path)
     w, h = image_path.size
     
     img_array = np.array(image_path)
@@ -37,9 +36,7 @@ def image_to_numpy(image_path):
 
 
 def visualize_molecule(smiles):
-    """
-    使用RDKit将SMILES转换为分子结构图
-    """
+
     try:
         mol = Chem.MolFromSmiles(smiles)
         if mol is None:
@@ -50,23 +47,16 @@ def visualize_molecule(smiles):
         return None
 
 def predict(input_image):
-    """
-    主要的推理函数
-    """
+
     try:
-        # 加载和初始化ONNX模型
         session = ort.InferenceSession("model.onnx")
         
-        # 预处理图片
-         # Example usage: #change thie image
         img_array,w,h = image_to_numpy(input_image)
         processed_image=np.expand_dims(img_array, 0)
-        
-        # 获取模型输入输出名称
+
         input_name = session.get_inputs()[0].name
         output_name = session.get_outputs()[0].name
-        
-        # 进行推理
+
         outputs = session.run(None, {input_name: processed_image})
         preds = {'pred_logits':torch.from_numpy(outputs[0]), 'pred_boxes':torch.from_numpy(outputs[1])}
         ori_size=torch.Tensor([w,h]).long().unsqueeze(0)
@@ -82,10 +72,6 @@ def predict(input_image):
             'scores': score_[selected_indices]
         }
         
-        # filtered_output_dict={image_path: output
-        # }
-        
-        
         x_center = (output["boxes"][:, 0] + output["boxes"][:, 2]) / 2
         y_center = (output["boxes"][:, 1] + output["boxes"][:, 3]) / 2
         center_coords = torch.stack((x_center, y_center), dim=1)
@@ -98,32 +84,30 @@ def predict(input_image):
                                                     bond_labels=bond_labels,  result=[])
         smiles, mol_rebuit = mol_from_graph_with_chiral(atoms_df, bonds_list)
         
-        # 使用RDKit生成分子结构图
         mol_image = visualize_molecule(smiles)
         
         if mol_image is None:
-            return "无效的SMILES字符串", None
+            return "Invalid SMILES", None
         
         return smiles, mol_image
         
     except Exception as e:
-        return f"发生错误: {str(e)}", None
+        return f"Error: {str(e)}", None
 
-# 创建Gradio界面
 iface = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="pil"),
+    inputs=gr.Image(label="Upload molecular image", type="pil", show_label=False).style(height=256),
     outputs=[
-        gr.Text(label="SMILES字符串"),
-        gr.Image(label="分子结构图")
+        gr.Image(label="Prediction").style(height=256),
+        gr.Text(label="SMILES").style(show_copy_button=True),
+        gr.Textbox(label="Molfile").style(show_copy_button=True),
     ],
-    title="化学结构OCR",
-    description="上传一张包含化学结构的图片，获取对应的SMILES表示和分子结构图。",
+    title="OCSR",
+    description="Convert a molecular image into SMILES.<br> ",
     examples=[
         ["example.png"]
     ]
 )
 
-# 启动应用
 if __name__ == "__main__":
     iface.launch()