Upload 2 files

app.py

@@ -1032,210 +1032,14 @@ def main():
                     import traceback
                     st.code(traceback.format_exc())
             
-            # Direct Download Annotated Image Button
+            # Automatic Annotation Info
             st.markdown("---")
-            st.markdown("### 📸 Download Annotated Structure Image")
-            
-            col1, col2 = st.columns([3, 1])
-            
-            with col1:
-                st.info("💡 Generate a structure visualization with RMSD and sequence information embedded")
-            
-            with col2:
-                if st.button("🖼️ Generate Annotated Image", use_container_width=True, type="primary"):
-                    with st.spinner("Generating annotated image..."):
-                        try:
-                            from visualization import extract_window_pdb, transform_pdb_string
-                            
-                            # Create a simple structure visualization using matplotlib
-                            # Since we can't capture py3Dmol directly, we'll create a matplotlib-based view
-                            
-                            # Extract structures
-                            ref_pdb = extract_window_pdb(
-                                selected_row['Ref_Path'],
-                                selected_row['Ref_Window']
-                            )
-                            
-                            query_pdb = extract_window_pdb(
-                                selected_row['Query_Path'],
-                                selected_row['Query_Window']
-                            )
-                            
-                            query_aligned_pdb = transform_pdb_string(
-                                query_pdb,
-                                selected_row['Rotation_Matrix'],
-                                selected_row['Query_COM'],
-                                selected_row['Ref_COM']
-                            )
-                            
-                            # Parse coordinates for visualization
-                            from rmsd_utils import parse_residue_atoms
-                            
-                            # Create a matplotlib-based 3D visualization
-                            import matplotlib.pyplot as plt
-                            from mpl_toolkits.mplot3d import Axes3D
-                            
-                            fig = plt.figure(figsize=(12, 9), dpi=150)
-                            ax = fig.add_subplot(111, projection='3d')
-                            
-                            # Function to extract coordinates from PDB string
-                            def get_coords_from_pdb_string(pdb_string):
-                                coords = []
-                                for line in pdb_string.split('\n'):
-                                    if line.startswith(('ATOM', 'HETATM')):
-                                        try:
-                                            x = float(line[30:38].strip())
-                                            y = float(line[38:46].strip())
-                                            z = float(line[46:54].strip())
-                                            atom_name = line[12:16].strip()
-                                            coords.append((x, y, z, atom_name))
-                                        except:
-                                            continue
-                                return coords
-                            
-                            # Get coordinates
-                            ref_coords = get_coords_from_pdb_string(ref_pdb)
-                            query_coords = get_coords_from_pdb_string(query_aligned_pdb)
-                            
-                            # Plot reference structure (blue)
-                            if ref_coords:
-                                ref_x = [c[0] for c in ref_coords]
-                                ref_y = [c[1] for c in ref_coords]
-                                ref_z = [c[2] for c in ref_coords]
-                                ax.scatter(ref_x, ref_y, ref_z, c='#4A90E2', s=40, alpha=0.8, label='Reference')
-                                
-                                # Connect backbone atoms
-                                backbone_atoms = ['P', "C4'", "C3'", "O3'"]
-                                ref_backbone = [(c[0], c[1], c[2]) for c in ref_coords if c[3] in backbone_atoms]
-                                if len(ref_backbone) > 1:
-                                    bb_x = [c[0] for c in ref_backbone]
-                                    bb_y = [c[1] for c in ref_backbone]
-                                    bb_z = [c[2] for c in ref_backbone]
-                                    ax.plot(bb_x, bb_y, bb_z, c='#4A90E2', linewidth=2, alpha=0.6)
-                            
-                            # Plot query structure (red)
-                            if query_coords:
-                                query_x = [c[0] for c in query_coords]
-                                query_y = [c[1] for c in query_coords]
-                                query_z = [c[2] for c in query_coords]
-                                ax.scatter(query_x, query_y, query_z, c='#E94B3C', s=40, alpha=0.8, label='Query (Aligned)')
-                                
-                                # Connect backbone atoms
-                                query_backbone = [(c[0], c[1], c[2]) for c in query_coords if c[3] in backbone_atoms]
-                                if len(query_backbone) > 1:
-                                    bb_x = [c[0] for c in query_backbone]
-                                    bb_y = [c[1] for c in query_backbone]
-                                    bb_z = [c[2] for c in query_backbone]
-                                    ax.plot(bb_x, bb_y, bb_z, c='#E94B3C', linewidth=2, alpha=0.6)
-                            
-                            # Set labels and title
-                            ax.set_xlabel('X (Å)', fontsize=10)
-                            ax.set_ylabel('Y (Å)', fontsize=10)
-                            ax.set_zlabel('Z (Å)', fontsize=10)
-                            ax.legend(fontsize=10, loc='upper right')
-                            
-                            # Set viewing angle
-                            ax.view_init(elev=20, azim=45)
-                            
-                            # Equal aspect ratio
-                            if ref_coords or query_coords:
-                                all_coords = ref_coords + query_coords
-                                all_x = [c[0] for c in all_coords]
-                                all_y = [c[1] for c in all_coords]
-                                all_z = [c[2] for c in all_coords]
-                                
-                                max_range = max(
-                                    max(all_x) - min(all_x),
-                                    max(all_y) - min(all_y),
-                                    max(all_z) - min(all_z)
-                                ) / 2.0
-                                
-                                mid_x = (max(all_x) + min(all_x)) / 2
-                                mid_y = (max(all_y) + min(all_y)) / 2
-                                mid_z = (max(all_z) + min(all_z)) / 2
-                                
-                                ax.set_xlim(mid_x - max_range, mid_x + max_range)
-                                ax.set_ylim(mid_y - max_range, mid_y + max_range)
-                                ax.set_zlim(mid_z - max_range, mid_z + max_range)
-                            
-                            plt.tight_layout()
-                            
-                            # Save to temporary buffer
-                            buf = io.BytesIO()
-                            plt.savefig(buf, format='png', dpi=150, bbox_inches='tight', facecolor='white')
-                            plt.close()
-                            buf.seek(0)
-                            
-                            # Now annotate this image
-                            annotated_img = annotate_alignment_image(
-                                image_data=buf.read(),
-                                rmsd=selected_row['RMSD'],
-                                ref_name=selected_row['Reference'],
-                                query_name=selected_row['Query'],
-                                ref_sequence=selected_row['Ref_Sequence'],
-                                query_sequence=selected_row['Query_Sequence'],
-                                output_format='JPEG'
-                            )
-                            
-                            # Generate filename
-                            ref_clean = selected_row['Reference'].replace('.pdb', '')
-                            query_clean = selected_row['Query'].replace('.pdb', '')
-                            filename = f"annotated_{ref_clean}_{query_clean}_RMSD_{selected_row['RMSD']:.3f}.jpg"
-                            
-                            # Show preview and download button
-                            st.success("✅ Annotated image generated!")
-                            st.image(annotated_img, caption="Annotated Structure Alignment", use_column_width=True)
-                            
-                            st.download_button(
-                                label="📥 Download Annotated JPEG",
-                                data=annotated_img.getvalue(),
-                                file_name=filename,
-                                mime="image/jpeg",
-                                use_container_width=True,
-                                help="Download JPEG with RMSD and sequence information"
-                            )
-                            
-                        except Exception as e:
-                            st.error(f"Error generating annotated image: {str(e)}")
-                            import traceback
-                            st.code(traceback.format_exc())
-                            
-                            # Fallback: offer the upload option
-                            st.info("💡 Alternatively, you can download a screenshot from the 3D viewer above using the '📷 Download PNG' button, then upload it below:")
-                            
-                            uploaded_screenshot = st.file_uploader(
-                                "Upload screenshot (PNG/JPG)",
-                                type=['png', 'jpg', 'jpeg'],
-                                key=f"screenshot_upload_fallback_{selected_viz_idx}"
-                            )
-                            
-                            if uploaded_screenshot is not None:
-                                try:
-                                    annotated_img = annotate_alignment_image(
-                                        image_data=uploaded_screenshot.read(),
-                                        rmsd=selected_row['RMSD'],
-                                        ref_name=selected_row['Reference'],
-                                        query_name=selected_row['Query'],
-                                        ref_sequence=selected_row['Ref_Sequence'],
-                                        query_sequence=selected_row['Query_Sequence'],
-                                        output_format='JPEG'
-                                    )
-                                    
-                                    st.image(annotated_img, use_column_width=True)
-                                    
-                                    ref_clean = selected_row['Reference'].replace('.pdb', '')
-                                    query_clean = selected_row['Query'].replace('.pdb', '')
-                                    filename = f"annotated_{ref_clean}_{query_clean}_RMSD_{selected_row['RMSD']:.3f}.jpg"
-                                    
-                                    st.download_button(
-                                        label="📥 Download Annotated JPEG",
-                                        data=annotated_img.getvalue(),
-                                        file_name=filename,
-                                        mime="image/jpeg",
-                                        use_container_width=True
-                                    )
-                                except Exception as e2:
-                                    st.error(f"Error annotating uploaded image: {str(e2)}")
+            st.success("✅ **Automatic Annotation:** When you click 'Download PNG' in the 3D viewer above, the image automatically includes RMSD, structure names, and sequences!")
+            st.info("💡 **Customize font size:** Use the 'Annotation Font Size' dropdown in the viewer controls (top-right) to choose from Small, Medium, Large (default), or Extra Large fonts!")
+
+
+
+
 
             # Show transformation details
             with st.expander("🔧 Transformation Details"):

visualization.py

@@ -315,6 +315,16 @@ def create_structure_visualization(ref_path, query_path, ref_window_indices, que
                     <option value="gray">Gray</option>
                 </select>
             </div>
+            
+            <div class="control-section">
+                <div class="section-title">Annotation Font Size</div>
+                <select id="annotationFontSize">
+                    <option value="small">Small (18pt/16pt/14pt)</option>
+                    <option value="medium">Medium (22pt/18pt/16pt)</option>
+                    <option value="large" selected>Large (28pt/22pt/18pt)</option>
+                    <option value="xlarge">Extra Large (36pt/28pt/22pt)</option>
+                </select>
+            </div>
         </div>
         
         <div class="rmsd-info">
@@ -607,14 +617,41 @@ def create_structure_visualization(ref_path, query_path, ref_window_indices, que
                     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';
+                    var refSeq = "{ref_sequence if ref_sequence else ''}";
+                    var querySeq = "{query_sequence if query_sequence else ''}";
+                    
+                    var filenameOriginal = 'alignment_' + refName + '_' + queryName + '_RMSD_' + rmsdValue + '.png';
+                    var filenameAnnotated = 'annotated_' + refName + '_' + queryName + '_RMSD_' + rmsdValue + '.png';
+                    
+                    // Get selected font size
+                    const fontSizeSelect = document.getElementById('annotationFontSize');
+                    const fontSizeOption = fontSizeSelect ? fontSizeSelect.value : 'large';
+                    
+                    // Define font sizes based on selection (all values are at 2x scale for high resolution)
+                    let fontSizes;
+                    switch(fontSizeOption) {{
+                        case 'small':
+                            fontSizes = {{ rmsd: 36, name: 32, seq: 28 }};  // 18pt/16pt/14pt at 2x
+                            break;
+                        case 'medium':
+                            fontSizes = {{ rmsd: 44, name: 36, seq: 32 }};  // 22pt/18pt/16pt at 2x
+                            break;
+                        case 'large':
+                            fontSizes = {{ rmsd: 56, name: 44, seq: 36 }};  // 28pt/22pt/18pt at 2x
+                            break;
+                        case 'xlarge':
+                            fontSizes = {{ rmsd: 72, name: 56, seq: 44 }};  // 36pt/28pt/22pt at 2x
+                            break;
+                        default:
+                            fontSizes = {{ rmsd: 56, name: 44, seq: 36 }};  // Default to large
+                    }}
                     
                     // Ensure viewer is rendered
                     if (viewer) {{
                         viewer.render();
                     }}
                     
-                    // Get the container element (includes canvas + all overlays)
+                    // Get the container element
                     const container = document.getElementById('container');
                     
                     if (!container) {{
@@ -630,34 +667,103 @@ def create_structure_visualization(ref_path, query_path, ref_window_indices, que
                         useCORS: true,
                         allowTaint: true
                     }}).then(function(canvas) {{
-                        // Convert to PNG
-                        const dataURL = canvas.toDataURL('image/png');
+                        // Create ANNOTATED version
+                        const annotatedCanvas = document.createElement('canvas');
+                        annotatedCanvas.width = canvas.width;
+                        annotatedCanvas.height = canvas.height;
+                        const ctx = annotatedCanvas.getContext('2d');
+                        
+                        // Draw the original image onto new canvas
+                        ctx.drawImage(canvas, 0, 0);
+                        
+                        // Add annotations
+                        const margin = 30;  // Scaled for 2x resolution
+                        const padding = 24;
+                        const lineSpacing = 16;
+                        
+                        // Prepare annotation text with selected font sizes
+                        const annotations = [
+                            {{ text: 'RMSD: ' + rmsdValue + ' Å', fontSize: fontSizes.rmsd, fontFamily: 'bold Arial', color: '#E94B3C' }},
+                            {{ text: '', fontSize: 20, fontFamily: 'Arial', color: '#333' }},  // Spacer
+                            {{ text: 'Reference: ' + refName, fontSize: fontSizes.name, fontFamily: 'Arial', color: '#333' }},
+                            {{ text: '  Seq: ' + refSeq, fontSize: fontSizes.seq, fontFamily: 'Courier New, monospace', color: '#666' }},
+                            {{ text: '', fontSize: 20, fontFamily: 'Arial', color: '#333' }},  // Spacer
+                            {{ text: 'Query: ' + queryName, fontSize: fontSizes.name, fontFamily: 'Arial', color: '#333' }},
+                            {{ text: '  Seq: ' + querySeq, fontSize: fontSizes.seq, fontFamily: 'Courier New, monospace', color: '#666' }}
+                        ];
+                        
+                        // Calculate box dimensions
+                        let maxWidth = 0;
+                        let totalHeight = padding * 2;
+                        const textMetrics = [];
                         
-                        // Create download link
-                        const link = document.createElement('a');
-                        link.download = filename;
-                        link.href = dataURL;
+                        annotations.forEach(ann => {{
+                            if (ann.text) {{
+                                ctx.font = ann.fontSize + 'px ' + ann.fontFamily;
+                                const metrics = ctx.measureText(ann.text);
+                                const height = ann.fontSize * 1.2;  // Approximate height
+                                textMetrics.push({{ width: metrics.width, height: height }});
+                                maxWidth = Math.max(maxWidth, metrics.width);
+                                totalHeight += height + lineSpacing;
+                            }} else {{
+                                textMetrics.push({{ width: 0, height: lineSpacing / 2 }});
+                                totalHeight += lineSpacing / 2;
+                            }}
+                        }});
+                        
+                        const boxWidth = maxWidth + padding * 2;
+                        const boxHeight = totalHeight;
+                        
+                        // Position box in bottom-left
+                        const boxX = margin;
+                        const boxY = annotatedCanvas.height - boxHeight - margin;
+                        
+                        // Draw semi-transparent white background with rounded corners
+                        ctx.fillStyle = 'rgba(255, 255, 255, 0.95)';
+                        const radius = 16;
+                        ctx.beginPath();
+                        ctx.moveTo(boxX + radius, boxY);
+                        ctx.lineTo(boxX + boxWidth - radius, boxY);
+                        ctx.quadraticCurveTo(boxX + boxWidth, boxY, boxX + boxWidth, boxY + radius);
+                        ctx.lineTo(boxX + boxWidth, boxY + boxHeight - radius);
+                        ctx.quadraticCurveTo(boxX + boxWidth, boxY + boxHeight, boxX + boxWidth - radius, boxY + boxHeight);
+                        ctx.lineTo(boxX + radius, boxY + boxHeight);
+                        ctx.quadraticCurveTo(boxX, boxY + boxHeight, boxX, boxY + boxHeight - radius);
+                        ctx.lineTo(boxX, boxY + radius);
+                        ctx.quadraticCurveTo(boxX, boxY, boxX + radius, boxY);
+                        ctx.closePath();
+                        ctx.fill();
+                        
+                        // Draw border
+                        ctx.strokeStyle = 'rgba(200, 200, 200, 0.95)';
+                        ctx.lineWidth = 2;
+                        ctx.stroke();
+                        
+                        // Draw text
+                        let currentY = boxY + padding;
+                        annotations.forEach((ann, idx) => {{
+                            if (ann.text) {{
+                                ctx.font = ann.fontSize + 'px ' + ann.fontFamily;
+                                ctx.fillStyle = ann.color;
+                                ctx.fillText(ann.text, boxX + padding, currentY + textMetrics[idx].height * 0.8);
+                                currentY += textMetrics[idx].height + lineSpacing;
+                            }} else {{
+                                currentY += textMetrics[idx].height;
+                            }}
+                        }});
+                        
+                        // Download ONLY the annotated PNG
+                        const annotatedDataURL = annotatedCanvas.toDataURL('image/png');
+                        const linkAnnotated = document.createElement('a');
+                        linkAnnotated.download = filenameAnnotated;
+                        linkAnnotated.href = annotatedDataURL;
+                        document.body.appendChild(linkAnnotated);
+                        linkAnnotated.click();
+                        document.body.removeChild(linkAnnotated);
                         
-                        // 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.');
-                        }}
+                        alert('Error creating images. Please try again.');
                     }});
                     
                 }} catch (error) {{