Spaces:
Sleeping
Sleeping
File size: 2,449 Bytes
3e76049 |
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 |
import numpy as np
import pydicom
import cv2
import tensorflow as tf
from tensorflow.keras.models import load_model
from model import create_simplified_model
import os
from PIL import Image, ImageDraw
def load_dicom_image(file_path):
"""Loads a DICOM image and converts it to a NumPy array."""
try:
ds = pydicom.dcmread(file_path)
image = ds.pixel_array
image = cv2.resize(image, (256, 256))
image = np.expand_dims(image, axis=-1)
image = image / 255.0
return image
except pydicom.errors.InvalidDicomError:
print(f"Error: Invalid DICOM file: {file_path}")
return None # Or raise a custom exception
except Exception as e: # added a catch all exception.
print(f"An unexpected error occurred: {e}")
return None
# Load or create model
try:
segmentation_model = load_model("pneumothorax_unet_model.h5")
print("Model loaded successfully!")
except Exception as e:
print(f"Error loading model: {e}")
def segment_pneumothorax(dicom_image_path):
"""Segments the pneumothorax region."""
image = load_dicom_image(dicom_image_path)
image = np.expand_dims(image, axis=0) # Add batch dimension
mask = segmentation_model.predict(image)[0]
mask = (mask > 0.5).astype(np.uint8) * 255 # Threshold and convert to 0-255
return mask
def overlay_mask(image, mask):
"""Overlays the mask on the image"""
image = np.squeeze(image) # remove channel dimension
image = (image * 255).astype(np.uint8) # convert back to 0-255.
image_color = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
mask_color = np.zeros_like(image_color)
mask_color[:, :, 0] = mask
overlay = cv2.addWeighted(image_color, 0.7, mask_color, 0.3, 0)
return overlay
def analyze_pneumothorax(dicom_file):
"""Combines loading, segmentation, and overlay."""
dicom_path = dicom_file.name
print(f"DICOM Path: {dicom_path}")
image = load_dicom_image(dicom_path)
if image is None:
return "Error: Invalid DICOM file uploaded."
mask = segment_pneumothorax(dicom_path)
overlayed_image = overlay_mask(image, mask)
print(f"Overlayed Image Type: {type(overlayed_image)}")
# Test Image
test_image = Image.new("RGB", (256, 256), "white")
draw = ImageDraw.Draw(test_image)
draw.ellipse((50, 50, 200, 200), fill="red")
return test_image |