| |
| from __future__ import annotations |
|
|
| import json |
| import os |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
| from PIL import Image |
|
|
| from flow_grpo.radiomics_reward import select_mask_roi |
|
|
|
|
| DEFAULT_TEST_JSONL = "/home/wenting/zr/wt_dataset/LIDC_IDRI/anno/cxr_synth_anno_mask_test.jsonl" |
|
|
|
|
| def _read_one_jsonl(path: Path) -> dict[str, Any]: |
| with open(path, "r", encoding="utf-8") as handle: |
| for line in handle: |
| if line.strip(): |
| return json.loads(line) |
| raise RuntimeError(f"No samples found in {path}") |
|
|
|
|
| def _load_mask_npz(path: Path, key: str = "mask") -> np.ndarray: |
| with np.load(path) as data: |
| selected_key = key if key in data.files else data.files[0] |
| return data[selected_key] |
|
|
|
|
| def _image_summary(path: Path) -> dict[str, Any]: |
| with Image.open(path) as image: |
| return {"path": str(path), "mode": image.mode, "size": list(image.size)} |
|
|
|
|
| def main() -> int: |
| jsonl_path = Path(os.environ.get("TEST_JSONL") or os.environ.get("TRAIN_JSONL") or DEFAULT_TEST_JSONL) |
| sample = _read_one_jsonl(jsonl_path) |
| output_image = Path(sample["output_image"]) |
| output_mask = Path(sample["output_mask"]) |
| mask = _load_mask_npz(output_mask, os.environ.get("MASK_NPZ_KEY", "mask")) |
|
|
| roi_lung, lung_channels = select_mask_roi(mask, mask_channels=[0, 1], mode="union") |
| roi_all, all_channels = select_mask_roi(mask, mask_channels="all", mode="union") |
| result = { |
| "jsonl": str(jsonl_path), |
| "sample_id": sample.get("sample_id"), |
| "patient_id": sample.get("patient_id"), |
| "input_images": sample.get("input_images"), |
| "gt_output_image": _image_summary(output_image), |
| "gt_output_mask": str(output_mask), |
| "mask_array_shape": list(mask.shape), |
| "mask_array_dtype": str(mask.dtype), |
| "channels_0_1": lung_channels, |
| "roi_pixels_channels_0_1": int(roi_lung.sum()), |
| "channels_all": all_channels, |
| "roi_pixels_all": int(roi_all.sum()), |
| "generated_mask_required": False, |
| } |
| print(json.dumps(result, indent=2, sort_keys=True)) |
| if int(roi_lung.sum()) <= 0: |
| raise SystemExit("Selected [0, 1] ROI is empty.") |
| if int(roi_all.sum()) <= 0: |
| raise SystemExit("Selected all-channel ROI is empty.") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
