Update app.py

app.py

@@ -5,10 +5,19 @@ from ONNX0630 import main as predict_smiles
 import tempfile
 import logging
 
-# 设置日志
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+def visualize_molecule(smiles):
+
+    try:
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None:
+            return None
+        img = Draw.MolToImage(mol)
+        return img
+    except:
+        return None
 
 def gradio_predict(image):
     try:
@@ -37,7 +46,7 @@ def gradio_predict(image):
         # 调用模型的 main 函数
         logger.info(f"传递给 main 的路径: {temp_path}")
         smiles = predict_smiles([temp_path])
-
+        mol_image = visualize_molecule(smiles)
         # 清理临时文件
         try:
             os.remove(temp_path)
@@ -45,7 +54,11 @@ def gradio_predict(image):
         except Exception as e:
             logger.warning(f"删除临时文件失败: {e}")
 
-        return smiles
+        if mol_image is None:
+            return None, smiles
+            
+        return mol_image, smiles
+        
     except FileNotFoundError as e:
         logger.error(f"文件路径错误: {e}")
         return f"错误: 文件路径问题 - {str(e)}"
@@ -63,7 +76,11 @@ def gradio_predict(image):
 iface = gr.Interface(
     fn=gradio_predict,
     inputs=gr.Image(type="pil"),
-    outputs=gr.Textbox(),
+    # outputs=gr.Textbox(),
+    outputs=[
+        gr.Image(label="Prediction"),
+        gr.Text(label="SMILES"),
+    ],
     title="化学结构转 SMILES",
     description="上传化学结构图片以获取其 SMILES 字符串"
 )