Radiuma_RnD/Dagster/Dagster_Op/dagster_radiuma/ops.py

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.")