Upload 2 files

app.py

@@ -44,6 +44,7 @@ st.set_page_config(
     initial_sidebar_state="expanded"
 )
 
+from image_annotator import annotate_alignment_image
 
 
 # Custom CSS - IMPROVED VERSION with larger fonts
@@ -497,27 +498,53 @@ def main():
     
     # Residue trimming controls - add early so they're available when needed
     st.sidebar.markdown("---")
-    st.sidebar.markdown("**🔧 Terminal Residue Trimming**")
+    st.sidebar.markdown("**🔧 Terminal Residue Trimming (Reference) **")
     col1, col2 = st.sidebar.columns(2)
     with col1:
-        n_term_trim = st.number_input(
-            "N-term trim",
+        n_term_trim_ref = st.number_input(
+            "N-term trim_ref",
             min_value=0,
             max_value=10,
             value=2,
             step=1,
             help="Number of residues to remove from 5' end",
-            key="n_term_trim"
+            key="n_term_trim_ref"
         )
     with col2:
-        c_term_trim = st.number_input(
-            "C-term trim",
+        c_term_trim_ref = st.number_input(
+            "C-term trim_ref",
             min_value=0,
             max_value=10,
             value=2,
             step=1,
             help="Number of residues to remove from 3' end",
-            key="c_term_trim"
+            key="c_term_trim_ref"
+        )
+    
+    
+    # Residue trimming controls - add early so they're available when needed
+    st.sidebar.markdown("---")
+    st.sidebar.markdown("**🔧 Terminal Residue Trimming (Query) **")
+    col1, col2 = st.sidebar.columns(2)
+    with col1:
+        n_term_trim_query = st.number_input(
+            "N-term trim_query",
+            min_value=0,
+            max_value=10,
+            value=2,
+            step=1,
+            help="Number of residues to remove from 5' end",
+            key="n_term_trim_query"
+        )
+    with col2:
+        c_term_trim_query = st.number_input(
+            "C-term trim_query",
+            min_value=0,
+            max_value=10,
+            value=2,
+            step=1,
+            help="Number of residues to remove from 3' end",
+            key="c_term_trim_query"
         )
     
     # Load structure data
@@ -624,8 +651,8 @@ def main():
                 
                 if selection_method == "Select by range":
                     current_selection = st.session_state['ref_selections'].get(selected_ref['name'], [])
-                    default_start = current_selection[0] + n_term_trim if current_selection else 3
-                    default_end = current_selection[-1] + 1 if current_selection else max(2, len(structure_info) - 2)
+                    default_start = current_selection[0] + n_term_trim_ref if current_selection else n_term_trim_ref
+                    default_end = current_selection[-1] + 1 if current_selection else max(n_term_trim_ref, len(structure_info) - c_term_trim_ref)
                     
                     c1, c2 = st.columns(2)
                     with c1:
@@ -653,13 +680,13 @@ def main():
                         
                 elif selection_method == "Select specific residues":
                     # Always use current trim values for default selection (updates when trim values change)
-                    default_names = [structure_info[i]['full_name'] for i in range(n_term_trim, len(structure_info)-c_term_trim)]
+                    default_names = [structure_info[i]['full_name'] for i in range(n_term_trim_ref, len(structure_info)-c_term_trim_ref)]
                     
                     selected_names = st.multiselect(
                         "Select residues",
                         options=[info['full_name'] for info in structure_info],
                         default=default_names,
-                        key=f"specific_ref_{selected_ref['name']}_n{n_term_trim}_c{c_term_trim}"
+                        key=f"specific_ref_{selected_ref['name']}_n{n_term_trim_ref}_c{c_term_trim_ref}"
                     )
                     
                     
@@ -724,8 +751,8 @@ def main():
                 
                 if selection_method == "Select by range":
                     current_selection = st.session_state['query_selections'].get(selected_query['name'], [])
-                    default_start = current_selection[0] + n_term_trim if current_selection else 3
-                    default_end = current_selection[-1] + 1 if current_selection else max(2, len(structure_info) - c_term_trim)
+                    default_start = current_selection[0] + n_term_trim_query if current_selection else 3
+                    default_end = current_selection[-1] + 1 if current_selection else max(2, len(structure_info) - c_term_trim_query)
                     
                     c1, c2 = st.columns(2)
                     with c1:
@@ -753,13 +780,13 @@ def main():
                         
                 elif selection_method == "Select specific residues":
                     # Always use current trim values for default selection (updates when trim values change)
-                    default_names = [structure_info[i]['full_name'] for i in range(n_term_trim, len(structure_info)-c_term_trim)]
+                    default_names = [structure_info[i]['full_name'] for i in range(n_term_trim_query, len(structure_info)-c_term_trim_query)]
                     
                     selected_names = st.multiselect(
                         "Select residues",
                         options=[info['full_name'] for info in structure_info],
                         default=default_names,
-                        key=f"specific_query_{selected_query['name']}_n{n_term_trim}_c{c_term_trim}"
+                        key=f"specific_query_{selected_query['name']}_n{n_term_trim_query}_c{c_term_trim_query}"
                     )
                     
                     name_to_idx = {info['full_name']: info['index'] for info in structure_info}
@@ -993,7 +1020,11 @@ def main():
                         selected_row['Rotation_Matrix'],
                         selected_row['Ref_COM'],
                         selected_row['Query_COM'],
-                        selected_row['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']
                     )
                     st.components.v1.html(viz_html, width=1400, height=750, scrolling=False)
                 except Exception as e:
@@ -1001,6 +1032,210 @@ def main():
                     import traceback
                     st.code(traceback.format_exc())
             
+            # Direct Download Annotated Image Button
+            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)}")
 
             # Show transformation details
             with st.expander("🔧 Transformation Details"):
@@ -1079,7 +1314,57 @@ def main():
                         help="Query structure aligned to reference"
                     )
                 
-                st.info("💡 **Tip:** Load reference and aligned query together in PyMOL/Chimera to examine the superposition")
+                # Combined aligned structure
+                st.markdown("---")
+                st.markdown("**Combined Aligned Structure (Reference + Query)**")
+                
+                # Create combined PDB with both structures
+                combined_pdb_lines = []
+                
+                # Add header information as REMARK records
+                combined_pdb_lines.append(f"REMARK Reference: {selected_row['Reference']}")
+                combined_pdb_lines.append(f"REMARK Reference Residues: {','.join(map(str, [i+1 for i in selected_row['Ref_Window']]))}")
+                combined_pdb_lines.append(f"REMARK Reference Sequence: {selected_row['Ref_Sequence']}")
+                combined_pdb_lines.append(f"REMARK Query: {selected_row['Query']}")
+                combined_pdb_lines.append(f"REMARK Query Residues: {','.join(map(str, [i+1 for i in selected_row['Query_Window']]))}")
+                combined_pdb_lines.append(f"REMARK Query Sequence: {selected_row['Query_Sequence']}")
+                combined_pdb_lines.append(f"REMARK RMSD: {selected_row['RMSD']:.3f} Angstroms")
+                combined_pdb_lines.append("MODEL        1")
+                
+                # Add reference atoms with chain A
+                for line in ref_pdb.split('\n'):
+                    if line.startswith(('ATOM', 'HETATM')):
+                        # Set chain to A for reference
+                        modified_line = line[:21] + 'A' + line[22:]
+                        combined_pdb_lines.append(modified_line)
+                
+                combined_pdb_lines.append("ENDMDL")
+                combined_pdb_lines.append("MODEL        2")
+                
+                # Add aligned query atoms with chain B
+                for line in query_aligned_pdb.split('\n'):
+                    if line.startswith(('ATOM', 'HETATM')):
+                        # Set chain to B for query
+                        modified_line = line[:21] + 'B' + line[22:]
+                        combined_pdb_lines.append(modified_line)
+                
+                combined_pdb_lines.append("ENDMDL")
+                combined_pdb_lines.append("END")
+                
+                combined_pdb = '\n'.join(combined_pdb_lines)
+                
+                combined_filename = f"aligned_{selected_row['Reference'].replace('.pdb', '')}_{selected_row['Query'].replace('.pdb', '')}_rmsd_{selected_row['RMSD']:.3f}.pdb"
+                
+                st.download_button(
+                    label="📥 Download Combined Aligned Structure",
+                    data=combined_pdb,
+                    file_name=combined_filename,
+                    mime="chemical/x-pdb",
+                    help="Reference (chain A) and aligned query (chain B) in one file",
+                    use_container_width=True
+                )
+                
+                st.info("💡 **Tip:** The combined PDB contains reference (chain A) and aligned query (chain B) - ready for PyMOL/Chimera")
         
         else:
             st.warning("No comparisons below RMSD threshold to visualize")

image_annotator.py

@@ -0,0 +1,221 @@
+"""
+Image Annotation Utility for RNA Structure Alignments
+Adds RMSD, reference/query names, and sequences directly to PNG/JPEG images
+"""
+
+from PIL import Image, ImageDraw, ImageFont
+import io
+
+
+def annotate_alignment_image(image_data, rmsd, ref_name, query_name, 
+                             ref_sequence=None, query_sequence=None, 
+                             output_format='PNG'):
+    """
+    Add text annotations to an alignment image.
+    
+    Args:
+        image_data: Either a file path (str), PIL Image object, or bytes
+        rmsd: RMSD value (float)
+        ref_name: Reference structure name (str)
+        query_name: Query structure name (str)
+        ref_sequence: Reference sequence (str, optional)
+        query_sequence: Query sequence (str, optional)
+        output_format: 'PNG' or 'JPEG'
+    
+    Returns:
+        BytesIO object containing the annotated image
+    """
+    
+    # Load the image
+    if isinstance(image_data, str):
+        img = Image.open(image_data)
+    elif isinstance(image_data, bytes):
+        img = Image.open(io.BytesIO(image_data))
+    elif isinstance(image_data, Image.Image):
+        img = image_data
+    else:
+        raise ValueError("image_data must be a file path, bytes, or PIL Image")
+    
+    # Convert to RGB if necessary (for JPEG compatibility)
+    if img.mode != 'RGB' and output_format == 'JPEG':
+        img = img.convert('RGB')
+    
+    # Create drawing context
+    draw = ImageDraw.Draw(img)
+    
+    # Try to use a better font, fall back to default if not available
+    try:
+        # Try to load a TrueType font
+        font_large = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 20)
+        font_medium = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 16)
+        font_small = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf", 14)
+    except:
+        # Fallback to default font
+        font_large = ImageFont.load_default()
+        font_medium = ImageFont.load_default()
+        font_small = ImageFont.load_default()
+    
+    # Get image dimensions
+    width, height = img.size
+    
+    # Define annotation box parameters
+    margin = 15
+    padding = 12
+    line_spacing = 8
+    
+    # Prepare text lines
+    lines = []
+    lines.append(("RMSD:", f"{rmsd:.3f} Å", font_large, True))  # Bold RMSD
+    lines.append(("", "", font_medium, False))  # Spacer
+    lines.append(("Reference:", ref_name, font_medium, False))
+    
+    if ref_sequence:
+        lines.append(("  Seq:", ref_sequence, font_small, False))
+    
+    lines.append(("", "", font_medium, False))  # Spacer
+    lines.append(("Query:", query_name, font_medium, False))
+    
+    if query_sequence:
+        lines.append(("  Seq:", query_sequence, font_small, False))
+    
+    # Calculate box dimensions
+    max_width = 0
+    total_height = padding * 2
+    
+    for label, value, font, is_bold in lines:
+        if label or value:
+            text = f"{label} {value}".strip()
+            bbox = draw.textbbox((0, 0), text, font=font)
+            text_width = bbox[2] - bbox[0]
+            text_height = bbox[3] - bbox[1]
+            max_width = max(max_width, text_width)
+            total_height += text_height + line_spacing
+        else:
+            total_height += line_spacing // 2
+    
+    box_width = max_width + padding * 2
+    box_height = total_height
+    
+    # Position box in bottom-left corner
+    box_x = margin
+    box_y = height - box_height - margin
+    
+    # Draw semi-transparent background box
+    overlay = Image.new('RGBA', img.size, (255, 255, 255, 0))
+    overlay_draw = ImageDraw.Draw(overlay)
+    
+    # White background with 95% opacity
+    overlay_draw.rounded_rectangle(
+        [(box_x, box_y), (box_x + box_width, box_y + box_height)],
+        radius=8,
+        fill=(255, 255, 255, 242)
+    )
+    
+    # Add subtle border
+    overlay_draw.rounded_rectangle(
+        [(box_x, box_y), (box_x + box_width, box_y + box_height)],
+        radius=8,
+        outline=(200, 200, 200, 242),
+        width=1
+    )
+    
+    # Composite the overlay
+    if img.mode == 'RGB':
+        img = img.convert('RGBA')
+    img = Image.alpha_composite(img, overlay)
+    
+    # Convert back to RGB if needed
+    if output_format == 'JPEG':
+        img = img.convert('RGB')
+    
+    # Recreate draw context after compositing
+    draw = ImageDraw.Draw(img)
+    
+    # Draw text
+    current_y = box_y + padding
+    
+    for label, value, font, is_bold in lines:
+        if not label and not value:
+            # Spacer
+            current_y += line_spacing // 2
+            continue
+        
+        text = f"{label} {value}".strip() if value else label
+        
+        # Choose color based on content
+        if "RMSD" in label:
+            color = (233, 75, 60)  # Red color for RMSD
+        elif label.startswith("  "):
+            color = (100, 100, 100)  # Gray for sequences
+        else:
+            color = (51, 51, 51)  # Dark gray for labels
+        
+        # Draw the text
+        draw.text((box_x + padding, current_y), text, fill=color, font=font)
+        
+        # Get text height for next line
+        bbox = draw.textbbox((0, 0), text, font=font)
+        text_height = bbox[3] - bbox[1]
+        current_y += text_height + line_spacing
+    
+    # Save to BytesIO
+    output = io.BytesIO()
+    img.save(output, format=output_format, quality=95 if output_format == 'JPEG' else None)
+    output.seek(0)
+    
+    return output
+
+
+def annotate_alignment_image_file(input_path, output_path, rmsd, ref_name, query_name,
+                                   ref_sequence=None, query_sequence=None):
+    """
+    Annotate an image file and save to a new file.
+    
+    Args:
+        input_path: Path to input image
+        output_path: Path to save annotated image
+        rmsd: RMSD value (float)
+        ref_name: Reference structure name (str)
+        query_name: Query structure name (str)
+        ref_sequence: Reference sequence (str, optional)
+        query_sequence: Query sequence (str, optional)
+    """
+    output_format = 'JPEG' if output_path.lower().endswith('.jpg') or output_path.lower().endswith('.jpeg') else 'PNG'
+    
+    annotated = annotate_alignment_image(
+        input_path, rmsd, ref_name, query_name, 
+        ref_sequence, query_sequence, output_format
+    )
+    
+    with open(output_path, 'wb') as f:
+        f.write(annotated.getvalue())
+
+
+# Example usage and testing
+if __name__ == "__main__":
+    # Example: Create a test image with annotations
+    from PIL import Image, ImageDraw
+    
+    # Create a simple test image
+    test_img = Image.new('RGB', (800, 600), color='white')
+    test_draw = ImageDraw.Draw(test_img)
+    
+    # Draw some simple shapes to simulate a structure
+    test_draw.ellipse([300, 200, 500, 400], fill='lightblue', outline='blue', width=3)
+    test_draw.ellipse([320, 220, 480, 380], fill='lightcoral', outline='red', width=3)
+    
+    # Annotate it
+    annotated = annotate_alignment_image(
+        test_img,
+        rmsd=1.234,
+        ref_name="6TNA_reference",
+        query_name="1EHZ_query",
+        ref_sequence="GCGGAU",
+        query_sequence="GCGGAU"
+    )
+    
+    # Save test image
+    with open('/tmp/test_annotated.png', 'wb') as f:
+        f.write(annotated.getvalue())
+    
+    print("Test image created: /tmp/test_annotated.png")