utils/visualization.py

import io
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image, ImageDraw, ImageFont


def draw_findings_on_image(image, findings):
    """
    Add annotations to X-ray image based on findings

    Args:
        image (PIL.Image): Original X-ray image
        findings (dict): Analysis findings with probabilities

    Returns:
        PIL.Image: Annotated image
    """
    # Create a copy of the image to draw on
    img = image.copy()
    draw = ImageDraw.Draw(img)

    # Get image dimensions
    width, height = img.size

    # Try to use a nice font, fall back to default if not available
    try:
        font = ImageFont.truetype("arial.ttf", 20)
        small_font = ImageFont.truetype("arial.ttf", 16)
    except IOError:
        font = ImageFont.load_default()
        small_font = ImageFont.load_default()

    # Add findings at the top
    y_position = 10
    for finding, probability in findings.items():
        if isinstance(probability, float):
            text = f"{finding}: {probability:.2f}"
            # Color code based on probability and finding type
            if finding == "No findings":
                color = (0, 128, 0)  # Green for no findings
            elif probability > 0.5:
                color = (255, 0, 0)  # Red for high probability issues
            else:
                color = (255, 165, 0)  # Orange for lower probability issues

            draw.text((10, y_position), text, fill=color, font=small_font)
            y_position += 25

    return img


def create_combined_visualization(image, image_results, text_results, combined_results):
    """
    Create a comprehensive visualization of all analysis results

    Args:
        image (PIL.Image): Original X-ray image
        image_results (dict): Image analysis results
        text_results (dict): Text analysis results
        combined_results (dict): Combined multimodal results

    Returns:
        PIL.Image: Visualization image
    """
    # Create a copy of the image to draw on
    img = image.copy()

    # Create a header with the recommendation
    recommendation = combined_results.get("Recommendation", "No recommendation")
    confidence = combined_results.get("Confidence", "N/A")

    # Create a white background for the header
    header_height = 60
    header_img = Image.new("RGB", (img.width, header_height), color=(255, 255, 255))
    header_draw = ImageDraw.Draw(header_img)

    # Try to use a nice font, fall back to default if not available
    try:
        font = ImageFont.truetype("arial.ttf", 18)
        small_font = ImageFont.truetype("arial.ttf", 14)
    except IOError:
        font = ImageFont.load_default()
        small_font = ImageFont.load_default()

    # Add recommendation text
    header_draw.text((10, 5), recommendation, fill=(0, 0, 0), font=font)
    header_draw.text(
        (10, 35), f"Confidence: {confidence}", fill=(100, 100, 100), font=small_font
    )

    # Combine the header and image
    combined_img = Image.new("RGB", (img.width, img.height + header_height))
    combined_img.paste(header_img, (0, 0))
    combined_img.paste(img, (0, header_height))

    return combined_img


def generate_report_plot(image_findings, text_findings):
    """
    Generate a comparison plot of image and text findings

    Args:
        image_findings (dict): Image analysis results
        text_findings (dict): Text analysis results

    Returns:
        bytes: PNG image data as bytes
    """
    # Create figure
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

    # Plot image findings
    image_labels = []
    image_values = []
    for k, v in image_findings.items():
        if isinstance(v, float):
            image_labels.append(k)
            image_values.append(v)

    # Sort by value for better visualization
    sorted_indices = np.argsort(image_values)[::-1]  # Descending order
    image_labels = [image_labels[i] for i in sorted_indices]
    image_values = [image_values[i] for i in sorted_indices]

    # Plot bars for image findings
    ax1.barh(image_labels, image_values, color="skyblue")
    ax1.set_xlim(0, 1)
    ax1.set_title("X-ray Analysis")
    ax1.set_xlabel("Probability")

    # Plot text findings (assuming text_findings has a structure to visualize)
    ax2.axis("off")  # Turn off axis
    ax2.text(0.1, 0.9, "Text Analysis Results:", fontweight="bold")
    y_pos = 0.8
    for key, value in text_findings.items():
        if key != "Entities":
            ax2.text(0.1, y_pos, f"{key}: {value}")
            y_pos -= 0.1

    # Adjust layout
    plt.tight_layout()

    # Convert to image bytes
    buf = io.BytesIO()
    plt.savefig(buf, format="png")
    buf.seek(0)

    # Close the plot to avoid memory leaks
    plt.close(fig)

    return buf.getvalue()