Add All Folders
This commit is contained in:
soorena62
@@ -0,0 +1,223 @@
|
||||
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.")
|
||||
Reference in New Issue
Block a user