Upload visualization.py

utils/visualization.py

@@ -0,0 +1,152 @@
+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()