image_annotator.py

"""
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")