soorena62
224 lines
6.9 KiB
Python
224 lines
6.9 KiB
Python
import os
|
|
import numpy as np
|
|
import pandas as pd
|
|
from pathlib import Path
|
|
import SimpleITK as sitk
|
|
from dagster import op, In, Out, DynamicOut, DynamicOutput
|
|
|
|
|
|
# Utility functions
|
|
def remove_ext(filename: str, exts):
|
|
fname = filename
|
|
for ext in exts:
|
|
if fname.lower().endswith(ext.lower()):
|
|
return fname[:-len(ext)]
|
|
return fname
|
|
|
|
def strip_suffixes(name_no_ext: str, suffixes):
|
|
for s in suffixes:
|
|
if name_no_ext.endswith(s):
|
|
return name_no_ext[:-len(s)]
|
|
return name_no_ext
|
|
|
|
def levenshtein(a: str, b: str) -> int:
|
|
if a == b:
|
|
return 0
|
|
m, n = len(a), len(b)
|
|
dp = list(range(n + 1))
|
|
for i in range(1, m + 1):
|
|
prev, dp[0] = dp[0], i
|
|
for j in range(1, n + 1):
|
|
cur = dp[j]
|
|
cost = 0 if a[i - 1] == b[j - 1] else 1
|
|
dp[j] = min(dp[j] + 1, dp[j - 1] + 1, prev + cost)
|
|
prev = cur
|
|
return dp[n]
|
|
|
|
def list_files(root: Path, exts):
|
|
files = []
|
|
for p in Path(root).rglob("*"):
|
|
if p.is_file() and any(str(p).lower().endswith(ext.lower()) for ext in exts):
|
|
files.append(p)
|
|
return sorted(files)
|
|
|
|
# Single-case readers
|
|
|
|
@op(
|
|
config_schema={
|
|
"image_path": str,
|
|
"mask_path": str,
|
|
},
|
|
out={"image": Out(), "mask": Out()}
|
|
)
|
|
def read_image_and_mask(context):
|
|
image_path = context.op_config["image_path"]
|
|
mask_path = context.op_config["mask_path"]
|
|
|
|
if not os.path.exists(image_path):
|
|
raise FileNotFoundError(f"Image not found: {image_path}")
|
|
if not os.path.exists(mask_path):
|
|
raise FileNotFoundError(f"Mask not found: {mask_path}")
|
|
|
|
image = sitk.ReadImage(image_path)
|
|
mask = sitk.ReadImage(mask_path)
|
|
img_arr = sitk.GetArrayFromImage(image)
|
|
msk_arr = sitk.GetArrayFromImage(mask)
|
|
|
|
context.log.info(f"[single] image:{img_arr.shape} mask:{msk_arr.shape}")
|
|
return img_arr, msk_arr
|
|
|
|
@op(
|
|
ins={"case": In()},
|
|
out={"image": Out(), "mask": Out()}
|
|
)
|
|
def read_image_and_mask_from_case(context, case: dict):
|
|
image_path = case["image_path"]
|
|
mask_path = case["mask_path"]
|
|
|
|
if not os.path.exists(image_path):
|
|
raise FileNotFoundError(f"Image not found: {image_path}")
|
|
if not os.path.exists(mask_path):
|
|
raise FileNotFoundError(f"Mask not found: {mask_path}")
|
|
|
|
image = sitk.ReadImage(image_path)
|
|
mask = sitk.ReadImage(mask_path)
|
|
img_arr = sitk.GetArrayFromImage(image)
|
|
msk_arr = sitk.GetArrayFromImage(mask)
|
|
|
|
context.log.info(f"[case] image:{img_arr.shape} mask:{msk_arr.shape}")
|
|
return img_arr, msk_arr
|
|
|
|
|
|
# Preprocess, features, writer
|
|
|
|
@op(out=Out())
|
|
def filter_image(context, image: np.ndarray):
|
|
context.log.info(f"[filter] image:{image.shape}")
|
|
return image
|
|
|
|
@op(ins={"image": In(), "mask": In()}, out=Out())
|
|
def extract_features(context, image: np.ndarray, mask: np.ndarray):
|
|
masked = image[mask > 0]
|
|
mean_val = float(masked.mean()) if masked.size > 0 else 0.0
|
|
features = {"mean_intensity": mean_val}
|
|
context.log.info(f"[features] {features}")
|
|
return features
|
|
|
|
|
|
RESULTS_DIR = Path("results")
|
|
RESULTS_DIR.mkdir(exist_ok=True)
|
|
|
|
@op(ins={"features": In(), "case_name": In()}, out=Out())
|
|
def write_outputs(context, features: dict, case_name: str):
|
|
# Define output CSV file path
|
|
csv_path = RESULTS_DIR / "features.csv"
|
|
|
|
# Prepare row with case name and features
|
|
row = {"case_name": case_name}
|
|
row.update(features)
|
|
|
|
# If CSV does not exist, create new file with header
|
|
if not csv_path.exists():
|
|
df = pd.DataFrame([row])
|
|
df.to_csv(csv_path, index=False)
|
|
else:
|
|
# Append new row to existing CSV
|
|
df = pd.DataFrame([row])
|
|
df.to_csv(csv_path, mode="a", header=False, index=False)
|
|
|
|
context.log.info(f"[write] features appended to {csv_path}")
|
|
|
|
# Optional Excel export (currently disabled)
|
|
# excel_path = RESULTS_DIR / "features.xlsx"
|
|
# df.to_excel(excel_path, index=False)
|
|
|
|
return str(csv_path)
|
|
|
|
|
|
# Batch case discovery
|
|
|
|
@op(
|
|
config_schema={
|
|
"images_dir": str,
|
|
"masks_dir": str,
|
|
"image_exts": list,
|
|
"mask_exts": list,
|
|
"mask_suffixes": list,
|
|
"allow_fuzzy_match": bool,
|
|
"fuzzy_max_distance": int,
|
|
"verbose": bool,
|
|
},
|
|
out=DynamicOut()
|
|
)
|
|
def enumerate_cases_auto(context):
|
|
cfg = context.op_config
|
|
images_dir = Path(cfg["images_dir"]).expanduser().resolve()
|
|
masks_dir = Path(cfg["masks_dir"]).expanduser().resolve()
|
|
image_exts = [ext.lower() for ext in cfg["image_exts"]]
|
|
mask_exts = [ext.lower() for ext in cfg["mask_exts"]]
|
|
mask_suffixes = cfg["mask_suffixes"]
|
|
allow_fuzzy = cfg["allow_fuzzy_match"]
|
|
max_dist = int(cfg["fuzzy_max_distance"])
|
|
verbose = bool(cfg.get("verbose", True))
|
|
|
|
img_files = list_files(images_dir, image_exts)
|
|
msk_files = list_files(masks_dir, mask_exts)
|
|
|
|
context.log.info(f"[discover] images_dir={images_dir} masks_dir={masks_dir}")
|
|
context.log.info(f"[discover] found images: {len(img_files)}")
|
|
for p in img_files:
|
|
context.log.info(f" - image: {p.name}")
|
|
context.log.info(f"[discover] found masks: {len(msk_files)}")
|
|
for p in msk_files:
|
|
context.log.info(f" - mask: {p.name}")
|
|
|
|
msk_index = {}
|
|
for m in msk_files:
|
|
base = strip_suffixes(remove_ext(m.name, mask_exts), mask_suffixes).lower()
|
|
msk_index.setdefault(base, []).append(m)
|
|
|
|
yielded = 0
|
|
skipped = 0
|
|
|
|
for img in img_files:
|
|
base = remove_ext(img.name, image_exts).lower()
|
|
candidates = msk_index.get(base, [])
|
|
chosen = None
|
|
reason = None
|
|
|
|
if candidates:
|
|
chosen = candidates[0]
|
|
else:
|
|
if allow_fuzzy and msk_files:
|
|
best = None
|
|
best_d = 10**9
|
|
for m in msk_files:
|
|
mbase = strip_suffixes(remove_ext(m.name, mask_exts), mask_suffixes).lower()
|
|
d = levenshtein(base, mbase)
|
|
if d < best_d:
|
|
best, best_d = m, d
|
|
if best is not None and best_d <= max_dist:
|
|
chosen = best
|
|
else:
|
|
reason = f"no exact or fuzzy match (best distance {best_d})"
|
|
|
|
if chosen is None:
|
|
skipped += 1
|
|
context.log.warning(f"[skip] image:{img.name} → no mask found ({reason}) | base(image)='{base}'")
|
|
continue
|
|
|
|
case = {
|
|
"image_path": str(img),
|
|
"mask_path": str(chosen),
|
|
"case_name": remove_ext(img.name, image_exts),
|
|
}
|
|
context.log.info(f"[pair] {case['case_name']} → image:{img.name} mask:{Path(chosen).name}")
|
|
yield DynamicOutput(case, mapping_key=case["case_name"])
|
|
yielded += 1
|
|
|
|
context.log.info(f"[summary] yielded={yielded} skipped={skipped}")
|
|
|
|
if yielded == 0:
|
|
context.log.error("No valid cases discovered. See [skip] messages above for reasons.")
|
|
raise Exception("No valid cases discovered. See logs for details.")
|