soorena62
109 lines
2.7 KiB
Python
109 lines
2.7 KiB
Python
import os
|
|
import numpy as np
|
|
import nibabel as nib
|
|
import pydicom
|
|
import nrrd
|
|
import cv2
|
|
|
|
from api.module import Module
|
|
from api.io_port import OutPort, NIFTIImageType
|
|
|
|
|
|
def load_nifti(path):
|
|
nii = nib.load(path)
|
|
return nii.get_fdata().astype(np.float32)
|
|
|
|
|
|
def load_dicom(path):
|
|
if os.path.isdir(path):
|
|
files = sorted([
|
|
os.path.join(path, f)
|
|
for f in os.listdir(path)
|
|
if not f.startswith(".")
|
|
])
|
|
slices = [pydicom.dcmread(f).pixel_array for f in files]
|
|
return np.stack(slices).astype(np.float32)
|
|
|
|
ds = pydicom.dcmread(path)
|
|
return ds.pixel_array.astype(np.float32)
|
|
|
|
|
|
def load_nrrd(path):
|
|
data, _ = nrrd.read(path)
|
|
return data.astype(np.float32)
|
|
|
|
|
|
def load_numpy(path):
|
|
return np.load(path).astype(np.float32)
|
|
|
|
def load_image(path):
|
|
img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
|
|
if img is None:
|
|
raise ValueError(f"Cannot read image file: {path}")
|
|
return img.astype(np.float32)
|
|
|
|
|
|
def load_any(path):
|
|
path = path.replace("\\", "/").lower()
|
|
|
|
if path.endswith((".nii", ".nii.gz")):
|
|
return load_nifti(path)
|
|
|
|
if path.endswith(".nrrd"):
|
|
return load_nrrd(path)
|
|
|
|
if path.endswith(".npy"):
|
|
return load_numpy(path)
|
|
|
|
if path.endswith((".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff")):
|
|
return load_image(path)
|
|
|
|
if os.path.isdir(path) or path.endswith(".dcm"):
|
|
return load_dicom(path)
|
|
|
|
raise ValueError(f"Unsupported image format: {path}")
|
|
|
|
|
|
def factory():
|
|
m = Module("ImageReader")
|
|
|
|
out_image = OutPort("image", NIFTIImageType(metadata={"role": "image"}))
|
|
out_mask = OutPort("mask", NIFTIImageType(metadata={"role": "mask"}))
|
|
|
|
m.add_out_port(out_image)
|
|
m.add_out_port(out_mask)
|
|
|
|
def execute(inputs):
|
|
image_dir = "data/images"
|
|
mask_dir = "data/masks"
|
|
|
|
image_files = [f for f in os.listdir(image_dir) if not f.startswith(".")]
|
|
mask_files = [f for f in os.listdir(mask_dir) if not f.startswith(".")]
|
|
|
|
if not image_files:
|
|
raise ValueError("No image found in data/images")
|
|
|
|
if not mask_files:
|
|
raise ValueError("No mask found in data/masks")
|
|
|
|
image_path = os.path.join(image_dir, image_files[0])
|
|
mask_path = os.path.join(mask_dir, mask_files[0])
|
|
|
|
try:
|
|
image = load_any(image_path)
|
|
except Exception as e:
|
|
raise ValueError(f"Cannot read image: {image_path} ({e})")
|
|
|
|
try:
|
|
mask = load_any(mask_path)
|
|
except Exception as e:
|
|
raise ValueError(f"Cannot read mask: {mask_path} ({e})")
|
|
|
|
return {
|
|
"image": image,
|
|
"mask": mask
|
|
}
|
|
|
|
m.execute = execute
|
|
return m
|