File size: 5,143 Bytes
1b052a1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 |
from pathlib import Path
from typing import Dict, Optional, Tuple
import numpy as np
from PIL import Image, ImageDraw
def create_fundus_overlay(
rgb_path: str,
av_path: Optional[str] = None,
disc_path: Optional[str] = None,
fovea_location: Optional[Tuple[int, int]] = None,
output_path: Optional[str] = None,
) -> np.ndarray:
"""
Create a visualization of a fundus image with overlaid segmentations and markers.
Args:
rgb_path: Path to the RGB fundus image
av_path: Optional path to artery-vein segmentation (1=artery, 2=vein, 3=intersection)
disc_path: Optional path to binary disc segmentation
fovea_location: Optional (x,y) tuple indicating the location of the fovea
output_path: Optional path to save the visualization image
Returns:
Numpy array containing the visualization image
"""
print(rgb_path, av_path, disc_path, fovea_location, output_path)
# Load RGB image
rgb_img = np.array(Image.open(rgb_path))
# Create output image starting with the RGB image
output_img = rgb_img.copy()
# Load and overlay AV segmentation if provided
if av_path:
av_mask = np.array(Image.open(av_path))
# Create masks for arteries (1), veins (2) and intersections (3)
artery_mask = av_mask == 1
vein_mask = av_mask == 2
intersection_mask = av_mask == 3
# Combine artery and intersection for visualization
artery_combined = np.logical_or(artery_mask, intersection_mask)
vein_combined = np.logical_or(vein_mask, intersection_mask)
# Apply colors: red for arteries, blue for veins
# Red channel - increase for arteries
output_img[artery_combined, 0] = 255
output_img[artery_combined, 1] = 0
output_img[artery_combined, 2] = 0
# Blue channel - increase for veins
output_img[vein_combined, 0] = 0
output_img[vein_combined, 1] = 0
output_img[vein_combined, 2] = 255
# Load and overlay optic disc segmentation if provided
if disc_path:
disc_mask = np.array(Image.open(disc_path)) > 0
# Apply white color for disc
output_img[disc_mask, :] = [255, 255, 255] # White
# Convert to PIL image for drawing the fovea marker
pil_img = Image.fromarray(output_img)
# Add fovea marker if provided
if fovea_location:
draw = ImageDraw.Draw(pil_img)
x, y = fovea_location
marker_size = (
min(pil_img.width, pil_img.height) // 50
) # Scale marker with image
# Draw yellow X at fovea location
draw.line(
[(x - marker_size, y - marker_size), (x + marker_size, y + marker_size)],
fill=(255, 255, 0),
width=2,
)
draw.line(
[(x - marker_size, y + marker_size), (x + marker_size, y - marker_size)],
fill=(255, 255, 0),
width=2,
)
# Convert back to numpy array
output_img = np.array(pil_img)
# Save output if path provided
if output_path:
Image.fromarray(output_img).save(output_path)
return output_img
def batch_create_overlays(
rgb_dir: Path,
output_dir: Path,
av_dir: Optional[Path] = None,
disc_dir: Optional[Path] = None,
fovea_data: Optional[Dict[str, Tuple[int, int]]] = None,
) -> None:
"""
Create visualization overlays for a batch of images.
Args:
rgb_dir: Directory containing RGB fundus images
output_dir: Directory to save visualization images
av_dir: Optional directory containing AV segmentations
disc_dir: Optional directory containing disc segmentations
fovea_data: Optional dictionary mapping image IDs to fovea coordinates
Returns:
List of paths to created visualization images
"""
# Create output directory if it doesn't exist
output_dir.mkdir(exist_ok=True, parents=True)
# Get all RGB images
rgb_files = list(rgb_dir.glob("*.png"))
if not rgb_files:
return []
# Process each image
for rgb_file in rgb_files:
image_id = rgb_file.stem
# Check for corresponding AV segmentation
av_file = None
if av_dir:
av_file_path = av_dir / f"{image_id}.png"
if av_file_path.exists():
av_file = str(av_file_path)
# Check for corresponding disc segmentation
disc_file = None
if disc_dir:
disc_file_path = disc_dir / f"{image_id}.png"
if disc_file_path.exists():
disc_file = str(disc_file_path)
# Get fovea location if available
fovea_location = None
if fovea_data and image_id in fovea_data:
fovea_location = fovea_data[image_id]
# Create output path
output_file = output_dir / f"{image_id}.png"
# Create and save overlay
create_fundus_overlay(
rgb_path=str(rgb_file),
av_path=av_file,
disc_path=disc_file,
fovea_location=fovea_location,
output_path=str(output_file),
)
|