File size: 8,551 Bytes
41c8683 | 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 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 | from PIL import Image
import numpy as np
from PIL import Image
from scipy.ndimage import binary_erosion, binary_dilation
from skimage.morphology import disk
def find_region(generated_image, erosion_dilation_radius=5):
red_channel = generated_image[:, :, 0] # red-channel
green_channel = generated_image[:, :, 1] # green-channel
blue_channel = generated_image[:, :, 2] # blue-channel
red_region = (red_channel > 100) & (green_channel < 80) & (blue_channel < 80)
selem = disk(erosion_dilation_radius)
mask = binary_erosion(red_region, structure=selem).astype(np.uint8)
mask = binary_dilation(mask, structure=selem).astype(np.uint8)
return mask
def define_organ_parts(mask):
if isinstance(mask, Image.Image):
mask = mask.convert("L")
mask = np.array(mask)
left_lung_coords = np.where(mask == 60)
right_lung_coords = np.where(mask == 120)
if left_lung_coords[0].size > 0:
left_min, left_max = left_lung_coords[0].min(), left_lung_coords[0].max()
left_lung_x_min, left_lung_x_max = left_lung_coords[1].min(), left_lung_coords[1].max()
left_upper_boundary = left_min + (left_max - left_min) // 3
left_middle_boundary = left_min + 2 * (left_max - left_min) // 3
else:
left_upper_boundary, left_middle_boundary, left_max = 0, 0, 0
left_lung_x_min, left_lung_x_max = 0, 0
if right_lung_coords[0].size > 0:
right_min, right_max = right_lung_coords[0].min(), right_lung_coords[0].max()
right_lung_x_min, right_lung_x_max = right_lung_coords[1].min(), right_lung_coords[1].max()
right_upper_boundary = right_min + (right_max - right_min) // 3
right_middle_boundary = right_min + 2 * (right_max - right_min) // 3
else:
right_upper_boundary, right_middle_boundary, right_max = 0, 0, 0
right_lung_x_min, right_lung_x_max = 0, 0
height, width = mask.shape[0], mask.shape[1]
organ_parts = {
"left upper lung": (
(np.arange(height)[:, None] <= left_upper_boundary) &
(np.arange(width) >= left_lung_x_min) &
(np.arange(width) <= left_lung_x_max)
),
"left middle lung": (
(np.arange(height)[:, None] > left_upper_boundary) &
(np.arange(height)[:, None] <= left_middle_boundary) &
(np.arange(width) >= left_lung_x_min) &
(np.arange(width) <= left_lung_x_max)
),
"left lower lung": (
(np.arange(height)[:, None] > left_middle_boundary) &
(np.arange(width) >= left_lung_x_min) &
(np.arange(width) <= left_lung_x_max)
),
"right upper lung": (
(np.arange(height)[:, None] <= right_upper_boundary) &
(np.arange(width) >= right_lung_x_min) &
(np.arange(width) <= right_lung_x_max)
),
"right middle lung": (
(np.arange(height)[:, None] > right_upper_boundary) &
(np.arange(height)[:, None] <= right_middle_boundary) &
(np.arange(width) >= right_lung_x_min) &
(np.arange(width) <= right_lung_x_max)
),
"right lower lung": (
(np.arange(height)[:, None] > right_middle_boundary) &
(np.arange(width) >= right_lung_x_min) &
(np.arange(width) <= right_lung_x_max)
),
"heart": (mask == 180),
"mediastinum": (mask == 240)
}
return organ_parts
def calculate_width(region_mask):
non_zero_columns = np.where(region_mask> 0)[1]
if len(non_zero_columns) == 0:
return 0
max_width = non_zero_columns.max() - non_zero_columns.min() + 1
return max_width
def process_organ_and_mask(disease, organ, mask):
organ_parts = define_organ_parts(organ)
if isinstance(organ, Image.Image):
organ = organ.convert("L")
organ = np.array(organ)
overlap_results = {}
for part, mask_part in organ_parts.items():
overlap_area = np.sum((mask_part > 0) & (mask > 0))
if overlap_area > 0:
overlap_results[part] = overlap_area
if overlap_results:
main_part = max(overlap_results, key=overlap_results.get)
if disease == 'Cardiomegaly':
if main_part == "heart":
location_label = main_part
organ_width = calculate_width(organ)
mask_width = calculate_width(mask)
cardio_ratio = mask_width / organ_width
if cardio_ratio <= 0.55:
severity = "mild"
elif 0.55 < cardio_ratio < 0.6:
severity = "moderate"
elif cardio_ratio >= 0.6:
severity = "severe"
else:
location_label = None
severity = None
if disease == 'Enlarged Cardiomediastinum':
organ_width = calculate_width(organ)
mask_width = calculate_width(mask)
if organ_width == 0 or mask_width == 0:
return None
ratio = mask_width / organ_width
location_label = "heart and mediastinum"
if ratio <= 0.55:
severity = "mild"
elif 0.55 < ratio < 0.6:
severity = "moderate"
else:
severity = "severe"
return disease, location_label, severity
else:
if overlap_results[main_part] > np.sum(mask > 0) * 0.7:
location_label = main_part
severity = "mild"
else:
left_regions = {"left upper lung", "left middle lung", "left lower lung"}
right_regions = {"right upper lung", "right middle lung", "right lower lung"}
active_regions = set(overlap_results.keys())
left_lung = (organ == 60)
left_overlap = active_regions & left_regions
left_lung_area = np.sum(left_lung)
left_overlap_area = np.sum((mask > 0) & (left_lung > 0))
left_lung_ratio = left_overlap_area / left_lung_area
right_lung = (organ == 120)
right_overlap = active_regions & right_regions
right_lung_area = np.sum(right_lung)
right_overlap_area = np.sum((mask > 0) & (right_lung > 0))
right_lung_ratio = right_overlap_area / right_lung_area
if left_overlap and right_overlap:
location_label = "biliteral lung"
elif left_overlap:
location_label = "left lung"
elif right_overlap:
location_label = "right lung"
else:
location_label = None
if disease == "Chest Tube" or disease == "Pacemaker":
severity = None
else:
if left_lung_ratio < 0.3 and right_lung_ratio < 0.3:
severity = "mild"
elif (
left_lung_ratio > 0.6 or
right_lung_ratio > 0.6 or
(left_lung_ratio + right_lung_ratio) > 0.6
):
severity = "severe"
else:
severity = "moderate"
generated_prompt = f"A Chest X-ray semantic mask with {severity} {disease} on {location_label}"
return generated_prompt
def post_process(generated_image, organ, disease, prompt):
mask = find_region(generated_image)
color_map = {
"Atelectasis": (255, 0, 0),
"Calcification": (0, 255, 0),
"Cardiomegaly": (0, 0, 255),
"Consolidation": (255, 255, 0),
"Diffuse Nodule": (255, 165, 0),
"Effusion": (0, 255, 255),
"Emphysema": (255, 0, 255),
"Fibrosis": (128, 0, 128),
"Fracture": (255, 192, 203),
"Mass": (173, 255, 47),
"Nodule": (0, 128, 255),
"Pleural Thickening": (75, 0, 130),
"Pneumothorax": (255, 105, 180)
}
generated_prompt = process_organ_and_mask(disease, organ, mask)
if generated_prompt == prompt:
organ_np = np.array(organ)
color = color_map.get(disease, [0, 0, 0])
organ_np[mask == 1] = color
return Image.fromarray(mask*255), Image.fromarray(organ_np)
else:
return None, None |