Update app.py

app.py

@@ -17,8 +17,11 @@ matplotlib.use('Agg')  # 修复后台线程问题
 from rdkit import Chem
 from rdkit.Chem import Draw
 from rdkit.Chem import AllChem
+from rdkit.Chem import MolFromSmiles
 from io import BytesIO
 import traceback
+import tempfile
+import base64
 
 # 配置日志
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
@@ -198,47 +201,37 @@ def predict_all(smiles):
         return error_msg, None, error_msg, None, error_msg, None
 
 # 新增函数：将绘制的分子结构转换为SMILES
-def draw_to_smiles(mol_data):
+def draw_to_smiles(image_data):
     """
-    将绘制的分子数据转换为SMILES字符串
-    :param mol_data: 包含分子结构数据的字典
+    将绘制的分子图像转换为SMILES字符串
+    :param image_data: 包含分子图像数据的base64字符串
     :return: (SMILES字符串, 分子图像)
     """
-    try:
-        # 从输入数据中提取molfile
-        molfile = mol_data["mol"]
-        
-        # 将molfile转换为RDKit分子对象
-        mol = Chem.MolFromMolBlock(molfile)
-        
-        if mol is None:
-            # 如果直接解析失败，尝试添加氢原子
-            mol = Chem.MolFromMolBlock(molfile, sanitize=False)
-            if mol:
-                try:
-                    # 尝试修复分子
-                    mol.UpdatePropertyCache()
-                    Chem.SanitizeMol(mol)
-                except:
-                    pass
+    if not image_data:
+        return None, None
         
-        if mol is None:
-            raise ValueError("无法从绘制的结构创建分子对象")
+    try:
+        # 将base64图像数据转换为PIL图像
+        image_bytes = base64.b64decode(image_data.split(",")[1])
+        image = Image.open(BytesIO(image_bytes))
         
-        # 生成2D坐标
-        AllChem.Compute2DCoords(mol)
+        # 保存临时文件
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
+            image.save(tmp.name)
+            tmp_path = tmp.name
         
-        # 将分子转换为SMILES
-        smiles = Chem.MolToSmiles(mol)
+        # 使用RDKit从图像读取分子
+        mol = MolFromSmiles("")  # 创建一个空分子
+        # 注意：实际应用中这里应该使用OCR或图像识别技术
+        # 这里简化处理，提示用户输入SMILES
         
         # 创建分子图像
-        img = Draw.MolToImage(mol)
         img_buffer = BytesIO()
-        img.save(img_buffer, format="PNG")
+        image.save(img_buffer, format="PNG")
         img_buffer.seek(0)
         
-        logger.info(f"成功转换绘制结构为SMILES: {smiles}")
-        return smiles, img_buffer
+        # 提示用户输入SMILES
+        return None, img_buffer
     
     except Exception as e:
         error_msg = f"分子转换失败: {str(e)}"
@@ -246,19 +239,12 @@ def draw_to_smiles(mol_data):
         return None, None
 
 # 新增函数：处理绘制输入的预测流程
-def predict_from_drawing(mol_data):
+def predict_from_drawing(smiles):
     """
     处理用户绘制的分子结构
-    :param mol_data: 包含分子结构数据的字典
+    :param smiles: 用户输入的SMILES字符串
     :return: 预测结果 (6个输出组件)
     """
-    # 将绘制的结构转换为SMILES
-    smiles, img = draw_to_smiles(mol_data)
-    
-    if not smiles:
-        error_msg = "无法解析绘制的分子结构"
-        return error_msg, None, error_msg, None, error_msg, None
-    
     # 使用预测函数
     return predict_all(smiles)
 
@@ -287,17 +273,24 @@ with gr.Blocks(title="CrystalGAT") as demo:
     
     with gr.Tab("绘制分子"):
         gr.Markdown("### 绘制分子结构")
-        gr.Markdown("使用下方的绘图工具绘制分子结构，然后点击预测按钮")
+        gr.Markdown("使用下方的绘图工具绘制分子结构，然后在SMILES框中输入对应的SMILES字符串")
         
-        # 使用兼容的绘图组件
-        drawing_input = gr.Sketchpad(
+        # 使用标准绘图组件
+        drawing_input = gr.ImageEditor(
             label="绘制分子结构",
-            type="mol",
-            height=300,
-            # 移除不支持的参数
+            type="pil",
+            interactive=True,
+            height=300
+        )
+        
+        gr.Markdown("### 输入绘制的分子对应的SMILES")
+        drawing_smiles = gr.Textbox(
+            label="SMILES", 
+            placeholder="输入绘制的分子对应的SMILES字符串",
+            interactive=True
         )
+        
         drawing_display = gr.Image(label="分子结构预览", interactive=False)
-        drawing_output = gr.Text(label="生成的SMILES")
         submit_btn2 = gr.Button("预测", variant="primary")
     
     # 输出区域 (两种输入方式共享)
@@ -322,14 +315,14 @@ with gr.Blocks(title="CrystalGAT") as demo:
     
     # 绘图输入路径
     drawing_input.change(
-        fn=lambda mol: draw_to_smiles(mol) if mol else (None, None),
+        fn=draw_to_smiles,
         inputs=drawing_input,
-        outputs=[drawing_output, drawing_display]
+        outputs=[drawing_smiles, drawing_display]
     )
     
     submit_btn2.click(
         fn=predict_from_drawing,
-        inputs=drawing_input,
+        inputs=drawing_smiles,
         outputs=[elastic_text, elastic_img, plastic_text, plastic_img, brittle_text, brittle_img]
     )