Spaces:

HouBioLab
/

MotifAlign

Sleeping

App Files Files Community

jiehou commited on Nov 19, 2025

Commit

8062cf7

verified ·

1 Parent(s): 91fbffc

Upload visualization.py

Browse files

Files changed (1) hide show

visualization.py +140 -4

visualization.py CHANGED Viewed

@@ -4,6 +4,7 @@ Uses py3Dmol for interactive molecular visualization
 """
 import numpy as np
 from rmsd_utils import (
     parse_residue_atoms,
     translate_rotate_coords,
@@ -13,7 +14,8 @@ from rmsd_utils import (
 def create_structure_visualization(ref_path, query_path, ref_window_indices, query_window_indices,
-                                   rotation_matrix, ref_com, query_com, rmsd=None):
     """
     Create an interactive 3D visualization of aligned structures.
@@ -26,11 +28,21 @@ def create_structure_visualization(ref_path, query_path, ref_window_indices, que
         ref_com: Center of mass of reference window
         query_com: Center of mass of query window
         rmsd: RMSD value (optional, for display)
     Returns:
         HTML string containing the py3Dmol visualization
     """
     # Read PDB files
     with open(ref_path) as f:
         ref_pdb = f.read()
@@ -84,6 +96,7 @@ def create_structure_visualization(ref_path, query_path, ref_window_indices, que
     <html>
     <head>
         <script src="https://3Dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
         <style>
             #container {{
                 width: 100%;
@@ -167,12 +180,25 @@ def create_structure_visualization(ref_path, query_path, ref_window_indices, que
                 bottom: 10px;
                 left: 10px;
                 background: rgba(255, 255, 255, 0.95);
-                padding: 10px 15px;
                 border-radius: 8px;
                 font-family: Arial, sans-serif;
                 font-size: 13px;
                 z-index: 1000;
                 box-shadow: 0 2px 8px rgba(0,0,0,0.1);
             }}
             .section-title {{
                 font-weight: bold;
@@ -181,6 +207,30 @@ def create_structure_visualization(ref_path, query_path, ref_window_indices, que
                 font-size: 12px;
                 text-transform: uppercase;
             }}
         </style>
     </head>
     <body>
@@ -268,7 +318,28 @@ def create_structure_visualization(ref_path, query_path, ref_window_indices, que
         </div>
         <div class="rmsd-info">
-            <strong>RMSD:</strong> <span style="color: #E94B3C; font-weight: bold;">{f"{rmsd:.3f}" if rmsd is not None else "N/A"} Å</span>
         </div>
         <script>
@@ -530,6 +601,71 @@ def create_structure_visualization(ref_path, query_path, ref_window_indices, que
                 }});
             }}
             // Initialize on load
             initViewer();
         </script>
@@ -670,4 +806,4 @@ def transform_pdb_string(pdb_string, rotation_matrix, query_com, ref_com=None):
         else:
             transformed_lines.append(line)
-    return '\n'.join(transformed_lines)

 """
 import numpy as np
+from pathlib import Path
 from rmsd_utils import (
     parse_residue_atoms,
     translate_rotate_coords,
 def create_structure_visualization(ref_path, query_path, ref_window_indices, query_window_indices,
+                                   rotation_matrix, ref_com, query_com, rmsd=None,
+                                   ref_name=None, query_name=None, ref_sequence=None, query_sequence=None):
     """
     Create an interactive 3D visualization of aligned structures.
         ref_com: Center of mass of reference window
         query_com: Center of mass of query window
         rmsd: RMSD value (optional, for display)
+        ref_name: Reference structure name (optional, for display)
+        query_name: Query structure name (optional, for display)
+        ref_sequence: Reference sequence (optional, for display)
+        query_sequence: Query sequence (optional, for display)
     Returns:
         HTML string containing the py3Dmol visualization
     """
+    # Extract simple names if not provided
+    if ref_name is None:
+        ref_name = Path(ref_path).stem
+    if query_name is None:
+        query_name = Path(query_path).stem
     # Read PDB files
     with open(ref_path) as f:
         ref_pdb = f.read()
     <html>
     <head>
         <script src="https://3Dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
+        <script src="https://cdnjs.cloudflare.com/ajax/libs/html2canvas/1.4.1/html2canvas.min.js"></script>
         <style>
             #container {{
                 width: 100%;
                 bottom: 10px;
                 left: 10px;
                 background: rgba(255, 255, 255, 0.95);
+                padding: 12px 15px;
                 border-radius: 8px;
                 font-family: Arial, sans-serif;
                 font-size: 13px;
                 z-index: 1000;
                 box-shadow: 0 2px 8px rgba(0,0,0,0.1);
+                max-width: 450px;
+            }}
+            .info-row {{
+                margin: 4px 0;
+                line-height: 1.4;
+            }}
+            .info-label {{
+                font-weight: bold;
+                color: #555;
+            }}
+            .info-value {{
+                color: #333;
+                font-family: 'Courier New', monospace;
             }}
             .section-title {{
                 font-weight: bold;
                 font-size: 12px;
                 text-transform: uppercase;
             }}
+            .download-section {{
+                position: absolute;
+                bottom: 10px;
+                right: 10px;
+                background: rgba(255, 255, 255, 0.95);
+                padding: 10px;
+                border-radius: 8px;
+                font-family: Arial, sans-serif;
+                z-index: 1000;
+                box-shadow: 0 2px 8px rgba(0,0,0,0.1);
+            }}
+            .download-btn {{
+                background: #4A90E2;
+                color: white;
+                border: none;
+                padding: 8px 16px;
+                border-radius: 4px;
+                cursor: pointer;
+                font-size: 13px;
+                font-weight: bold;
+            }}
+            .download-btn:hover {{
+                background: #357ABD;
+            }}
         </style>
     </head>
     <body>
         </div>
         <div class="rmsd-info">
+            <div class="info-row">
+                <span class="info-label">RMSD:</span>
+                <span style="color: #E94B3C; font-weight: bold; font-size: 14px;">{f"{rmsd:.3f}" if rmsd is not None else "N/A"} Å</span>
+            </div>
+            <div style="margin-top: 8px; padding-top: 8px; border-top: 1px solid #ddd;">
+                <div class="info-row">
+                    <span class="info-label">Reference:</span>
+                    <span class="info-value">{ref_name}</span>
+                </div>
+                {f'<div class="info-row" style="margin-left: 15px; font-size: 12px;"><span class="info-label">Seq:</span> <span class="info-value">{ref_sequence}</span></div>' if ref_sequence else ''}
+            </div>
+            <div style="margin-top: 6px;">
+                <div class="info-row">
+                    <span class="info-label">Query:</span>
+                    <span class="info-value">{query_name}</span>
+                </div>
+                {f'<div class="info-row" style="margin-left: 15px; font-size: 12px;"><span class="info-label">Seq:</span> <span class="info-value">{query_sequence}</span></div>' if query_sequence else ''}
+            </div>
+        </div>
+        <div class="download-section">
+            <button class="download-btn" onclick="downloadImage()">📷 Download PNG</button>
         </div>
         <script>
                 }});
             }}
+            function downloadImage() {{
+                try {{
+                    // Generate filename with metadata
+                    var refName = "{ref_name}".replace('.pdb', '');
+                    var queryName = "{query_name}".replace('.pdb', '');
+                    var rmsdValue = "{f'{rmsd:.3f}' if rmsd is not None else 'NA'}";
+                    var filename = 'alignment_' + refName + '_' + queryName + '_RMSD_' + rmsdValue + '.png';
+                    // Ensure viewer is rendered
+                    if (viewer) {{
+                        viewer.render();
+                    }}
+                    // Get the container element (includes canvas + all overlays)
+                    const container = document.getElementById('container');
+                    if (!container) {{
+                        alert('Container not ready. Please wait and try again.');
+                        return;
+                    }}
+                    // Use html2canvas to capture the entire container with overlays
+                    html2canvas(container, {{
+                        backgroundColor: '#ffffff',
+                        scale: 2,  // Higher resolution
+                        logging: false,
+                        useCORS: true,
+                        allowTaint: true
+                    }}).then(function(canvas) {{
+                        // Convert to PNG
+                        const dataURL = canvas.toDataURL('image/png');
+                        // Create download link
+                        const link = document.createElement('a');
+                        link.download = filename;
+                        link.href = dataURL;
+                        // Trigger download
+                        document.body.appendChild(link);
+                        link.click();
+                        document.body.removeChild(link);
+                    }}).catch(function(error) {{
+                        console.error('html2canvas error:', error);
+                        // Fallback: just capture the canvas without overlays
+                        const canvas = document.querySelector('#container canvas');
+                        if (canvas) {{
+                            const dataURL = canvas.toDataURL('image/png');
+                            const link = document.createElement('a');
+                            link.download = filename;
+                            link.href = dataURL;
+                            document.body.appendChild(link);
+                            link.click();
+                            document.body.removeChild(link);
+                        }} else {{
+                            alert('Download failed. Please try again.');
+                        }}
+                    }});
+                }} catch (error) {{
+                    console.error('PNG download error:', error);
+                    alert('Error downloading PNG: ' + error.message);
+                }}
+            }}
             // Initialize on load
             initViewer();
         </script>
         else:
             transformed_lines.append(line)
+    return '\n'.join(transformed_lines)