Add interactive 3D molecule viewer with 3Dmol.js

app.py

@@ -80,31 +80,27 @@ def name_to_3d_molecule(name: str) -> tuple:
         # Import needed modules
         import io
         import base64
+        import json
         from rdkit.Chem.Draw import rdMolDraw2D
         
         # Generate 2D image for preview
         mol_2d = Chem.MolFromSmiles(smiles)  # Get 2D version
         
         # Create drawer for 2D structure
-        drawer = rdMolDraw2D.MolDraw2DCairo(600, 600)
+        drawer = rdMolDraw2D.MolDraw2DCairo(500, 500)
         drawer.DrawMolecule(mol_2d)
         drawer.FinishDrawing()
         img_2d_data = drawer.GetDrawingText()
         img_2d_str = base64.b64encode(img_2d_data).decode()
         
-        # Create 3D structure image with atom indices
-        drawer_3d = rdMolDraw2D.MolDraw2DCairo(600, 600)
-        drawer_3d.drawOptions().addAtomIndices = True
-        AllChem.Compute2DCoords(mol)
-        drawer_3d.DrawMolecule(mol)
-        drawer_3d.FinishDrawing()
-        img_3d_data = drawer_3d.GetDrawingText()
-        img_3d_str = base64.b64encode(img_3d_data).decode()
-        
         # Get SMILES string
         canonical_smiles = Chem.MolToSmiles(mol_2d)
+        
+        # Escape SDF for JavaScript
+        sdf_json = json.dumps(sdf_content)
+        viewer_id = f"viewer_{abs(hash(name))}"
 
-        # Create HTML with both views
+        # Create HTML with 3Dmol.js viewer
         html_content = f"""
         <div style="padding: 20px;">
             <h3 style="text-align: center;">{name}</h3>
@@ -113,20 +109,65 @@ def name_to_3d_molecule(name: str) -> tuple:
             <div style="display: flex; gap: 20px; justify-content: center; flex-wrap: wrap; margin-top: 20px;">
                 <div style="text-align: center;">
                     <h4>2D Structure</h4>
-                    <img src="data:image/png;base64,{img_2d_str}" style="max-width: 600px; width: 100%; border: 1px solid #ddd; border-radius: 8px; padding: 10px; background: white;">
+                    <img src="data:image/png;base64,{img_2d_str}" style="max-width: 500px; width: 100%; border: 1px solid #ddd; border-radius: 8px; padding: 10px; background: white;">
                 </div>
                 
                 <div style="text-align: center;">
-                    <h4>3D Structure with Atom Labels</h4>
-                    <img src="data:image/png;base64,{img_3d_str}" style="max-width: 600px; width: 100%; border: 1px solid #ddd; border-radius: 8px; padding: 10px; background: white;">
-                    <p style="margin-top: 10px; font-size: 12px; color: #666;">3D optimized structure with atom indices</p>
+                    <h4>3D Interactive Viewer</h4>
+                    <div id="{viewer_id}" style="width: 500px; height: 500px; border: 1px solid #ddd; border-radius: 8px; background: white; position: relative;"></div>
                 </div>
             </div>
-            
-            <div style="margin-top: 20px; padding: 15px; background-color: #e7f3ff; border: 2px solid #2196F3; border-radius: 8px; max-width: 800px; margin-left: auto; margin-right: auto;">
-                <p style="color: #0d47a1; margin: 0;"><strong>💡 Tip:</strong> The SDF content below contains full 3D coordinates. Copy and paste into molecular viewers like PyMOL, Avogadro, or ChimeraX for interactive 3D visualization.</p>
-            </div>
         </div>
+        
+        <script src="https://code.jquery.com/jquery-3.6.0.min.js"></script>
+        <script src="https://3dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
+        <script>
+            (function() {{
+                var sdfData = {sdf_json};
+                var viewerId = "{viewer_id}";
+                
+                function init3DMol() {{
+                    try {{
+                        var element = document.getElementById(viewerId);
+                        if (!element) {{
+                            console.error('Viewer element not found:', viewerId);
+                            return;
+                        }}
+                        
+                        var config = {{ backgroundColor: 'white' }};
+                        var viewer = $3Dmol.createViewer(element, config);
+                        
+                        viewer.addModel(sdfData, "sdf");
+                        viewer.setStyle({{}}, {{
+                            stick: {{ radius: 0.15, colorscheme: 'Jmol' }},
+                            sphere: {{ radius: 0.4, colorscheme: 'Jmol' }}
+                        }});
+                        
+                        viewer.addLabels({{ elem: true }}, {{
+                            fontSize: 12,
+                            fontColor: 'black',
+                            backgroundColor: 'white',
+                            backgroundOpacity: 0.7
+                        }});
+                        
+                        viewer.zoomTo();
+                        viewer.render();
+                        viewer.zoom(1.2);
+                        
+                        console.log('3DMol viewer initialized successfully');
+                    }} catch (e) {{
+                        console.error('Error initializing 3DMol viewer:', e);
+                    }}
+                }}
+                
+                // Wait for both jQuery and 3Dmol to load
+                if (typeof $3Dmol !== 'undefined') {{
+                    init3DMol();
+                }} else {{
+                    setTimeout(init3DMol, 1000);
+                }}
+            }})();
+        </script>
         """
 
         return html_content, sdf_content
@@ -180,11 +221,11 @@ molecule_3d_interface = gr.Interface(
     fn=name_to_3d_molecule,
     inputs=gr.Textbox(label="Chemical Name", placeholder="e.g., benzene, aspirin, caffeine, glucose"),
     outputs=[
-        gr.HTML(label="2D and 3D Molecular Structures"),
-        gr.Textbox(label="3D SDF Content (for PyMOL, Avogadro, ChimeraX)", lines=10, max_lines=20, visible=True)
+        gr.HTML(label="2D and 3D Interactive Molecular Viewer"),
+        gr.Textbox(label="3D SDF Content (Optional - for external viewers)", lines=10, max_lines=20, visible=False)
     ],
     api_name="name_to_molecule",
-    description="View 2D structure and 3D optimized structure with atom labels. Copy the SDF content to view in external molecular viewers.",
+    description="View 2D structure and interactive 3D molecule viewer with atom labels. Rotate and zoom the 3D structure!",
     cache_examples=False,
 )