Update main.py

main.py

@@ -3,11 +3,14 @@ from pydantic import BaseModel
 import subprocess
 from typing import List
 from fragment_processor import fragmentize_molecule
-
+from fastapi.responses import StreamingResponse
 import torch
 import pandas as pd
 from rdkit import Chem
-from rdkit.Chem import Descriptors, QED
+from rdkit.Chem import Descriptors, QED, Draw
+from io import BytesIO
+from PIL import Image
+import io
 from generate import GenerateRunner
 from dataset import Dataset
 from combine_mol import connect_constVar_try
@@ -147,6 +150,33 @@ def run_generate_runner(const_smiles, var_smiles, main_cls, minor_cls, delta_val
     return result
 
 
+@app.get("/smiles2img")
+async def smiles2img(smiles: str = Query(..., description="SMILES string of the molecule")):
+    
+    # 生成分子对象
+    print("---开始生成分子图像--")
+    print(smiles)
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is None:
+        return {"error": "Invalid SMILES string"}
+
+    # 创建绘图对象
+    drawer = Draw.MolDraw2DCairo(200, 200)
+    drawer.SetFontSize(1.0)
+    drawer.DrawMolecule(mol)
+    drawer.FinishDrawing()
+
+    # 将绘制的图像转换为PIL图像对象
+    pil_image = Image.open(io.BytesIO(drawer.GetDrawingText()))
+
+    # 创建字节流
+    strIO = BytesIO()
+    pil_image.save(strIO, "PNG")
+    strIO.seek(0)
+
+    # 返回图像作为流
+    return StreamingResponse(strIO, media_type="image/png")
+
 @app.get("/fragmentize", response_model=FragmentResponse)
 async def fragmentize(smiles: str = Query(..., description="SMILES string of the molecule")):
     try: