fix: enabled 3d viewer

Steinbeck-Lab · Mar 20, 2023 · 043d4ad · 043d4ad
1 parent 7268750
commit 043d4ad
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 5 deletions.
diff --git a/app/routers/chem.py b/app/routers/chem.py
@@ -12,7 +12,9 @@
 from app.modules.classyfire import classify, result
 from app.modules.cdkmodules import getCDKSDGMol
 from app.modules.depict import getRDKitDepiction, getCDKDepiction
-from app.modules.depict3D import get3DDepiction
+from app.modules.depict3D import get3Dconformers
+
+from fastapi.templating import Jinja2Templates
 
 router = APIRouter(
     prefix="/chem",
@@ -21,6 +23,7 @@
     responses={404: {"description": "Not found"}},
 )
 
+templates = Jinja2Templates(directory="app/templates")
 
 @router.get("/")
 async def chem_index():
@@ -129,16 +132,16 @@ async def depict_molecule(
             )
 
 
-@router.get("/depict3D")
+@router.get("/depict3D", response_class=HTMLResponse)
 async def depict3D_molecule(
+    request: Request,
     smiles: str,
     width: Optional[int] = 512,
     height: Optional[int] = 512,
     style: Optional[str] = "stick",
 ):
     if smiles:
-        content__ = get3DDepiction(smiles, [width, height], style)
-        return HTMLResponse(content=content__.render(), status_code=200)
+        return templates.TemplateResponse("mol.html", {"request": request, "molecule": Chem.MolToMolBlock(get3Dconformers(smiles))})
 
 
 # @app.get("/molecules/", response_model=List[schemas.Molecule])

diff --git a/app/templates/mol.html b/app/templates/mol.html
@@ -0,0 +1,27 @@
+<html>
+<head>
+<title>3D Molecule Viewer</title>
+<script src="https://code.jquery.com/jquery-3.6.3.min.js"></script>
+<script src="https://cdnjs.cloudflare.com/ajax/libs/3Dmol/2.0.1/3Dmol.js"></script>
+<style>
+head, body {
+	margin: 0;
+	border: 0;
+	padding: 0;
+	max-height: 100vh
+}
+</style>
+<script>
+$(document).ready(function() {
+    var viewer = $3Dmol.createViewer("viewer");
+    viewer.setBackgroundColor(0xffffff);
+    viewer.addModel(`{{ molecule }}`, "sdf");
+    viewer.setStyle({stick:{}});
+    viewer.zoomTo();
+    viewer.render();
+});
+</script>
+<body>
+	<div id="viewer" style="width: 100%; height: 100vh; margin: 0; padding: 0; border: 0;"></div>
+</body>
+</html>
diff --git a/requirements.txt b/requirements.txt
@@ -16,4 +16,5 @@ efficientnet
 tensorflow==2.10.0
 pyheif==0.7.1
 selfies>=2.1.1
-py3Dmol
+py3Dmol
+jinja2