Create app.py

app.py

@@ -0,0 +1,88 @@
+import gradio as gr
+import onnxruntime as ort
+import numpy as np
+from PIL import Image
+import io
+import rdkit
+from rdkit import Chem
+from rdkit.Chem import Draw
+
+def preprocess_image(image):
+    """
+    预处理输入图片
+    """
+    # 将图片调整为模型所需的输入尺寸
+    image = image.resize((224, 224))  # 根据实际模型需求调整尺寸
+    # 转换为numpy数组并归一化
+    img_array = np.array(image)
+    img_array = img_array.astype(np.float32) / 255.0
+    # 添加批次维度
+    img_array = np.expand_dims(img_array, axis=0)
+    # 根据模型训练时的预处理方式进行调整
+    img_array = img_array.transpose(0, 3, 1, 2)  # BHWC to BCHW
+    return img_array
+
+def visualize_molecule(smiles):
+    """
+    使用RDKit将SMILES转换为分子结构图
+    """
+    try:
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None:
+            return None
+        img = Draw.MolToImage(mol)
+        return img
+    except:
+        return None
+
+def predict(input_image):
+    """
+    主要的推理函数
+    """
+    try:
+        # 加载和初始化ONNX模型
+        session = ort.InferenceSession("model.onnx")  # 替换为实际模型路径
+        
+        # 预处理图片
+        processed_image = preprocess_image(input_image)
+        
+        # 获取模型输入输出名称
+        input_name = session.get_inputs()[0].name
+        output_name = session.get_outputs()[0].name
+        
+        # 进行推理
+        predictions = session.run([output_name], {input_name: processed_image})
+        
+        # 假设模型输出是SMILES字符串
+        smiles = predictions[0][0]  # 根据实际模型输出格式调整
+        
+        # 使用RDKit生成分子结构图
+        mol_image = visualize_molecule(smiles)
+        
+        if mol_image is None:
+            return "无效的SMILES字符串", None
+        
+        return smiles, mol_image
+        
+    except Exception as e:
+        return f"发生错误: {str(e)}", None
+
+# 创建Gradio界面
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=[
+        gr.Text(label="SMILES字符串"),
+        gr.Image(label="分子结构图")
+    ],
+    title="化学结构OCR",
+    description="上传一张包含化学结构的图片，获取对应的SMILES表示和分子结构图。",
+    examples=[
+        ["example1.jpg"],  # 添加示例图片
+        ["example2.jpg"]
+    ]
+)
+
+# 启动应用
+if __name__ == "__main__":
+    iface.launch()