Add All Folders

Dagster/API_R&D/modules/image_reader.py

@@ -0,0 +1,108 @@
+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