""" Per-model inference entry points for the 6 AutoMSC_challenge baselines. For each baseline model in /mnt/pool/datasets/CY/AutoMSC_challenge/baseline_models/ this file exposes a function that: - loads the fold_0 `checkpoint_best.pth` - runs segmentation + classification on one case - writes the segmentation mask (NIfTI) to disk - writes an mp4 video overlay of the first modality + segmentation - returns (seg_mask_path, overlay_video_path, cls_results_dict) The cls_results_dict maps the human-readable class name (from dataset.json["classification_labels"]) to its predicted probability. """ import os import gc import json import shutil import subprocess import itertools from typing import List, Sequence, Tuple, Union, Optional import numpy as np import torch import cv2 import SimpleITK as sitk from tqdm import tqdm import nnunetv2 from batchgenerators.utilities.file_and_folder_operations import load_json, join from nnunetv2.utilities.plans_handling.plans_handler import PlansManager, ConfigurationManager from nnunetv2.utilities.label_handling.label_handling import LabelManager, determine_num_input_channels from nnunetv2.utilities.find_class_by_name import recursive_find_python_class from nnunetv2.inference.sliding_window_prediction import compute_gaussian from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor from nnunetv2.utilities.helpers import empty_cache from nnunetv2.imageio.simpleitk_reader_writer import SimpleITKIO from nnunetv2.preprocessing.resampling.default_resampling import fast_resample_logit_to_shape from acvl_utils.cropping_and_padding.padding import pad_nd_image from acvl_utils.cropping_and_padding.bounding_boxes import bounding_box_to_slice # ─── Registry of baseline models ──────────────────────────────────────────────── BASELINE_ROOT = "./baseline_models" RAW_ROOT = "." MODEL_REGISTRY = { "Dataset001_PETWB_Lung": { "plans_folder": "nnUNetCLSTrainerMTL__nnUNetPlans__3d_fullres", "dataset_json_fallback": f"{RAW_ROOT}/Dataset001_PETWB_Lung/dataset.json", }, "Dataset002_BMLMPS_FLAIR": { "plans_folder": "nnUNetCLSTrainerMTL__nnUNetPlans__3d_fullres", "dataset_json_fallback": f"{RAW_ROOT}/Dataset002_BMLMPS_FLAIR/dataset.json", }, "Dataset003_BMLMPS_T1CE": { "plans_folder": "nnUNetCLSTrainerMTL__nnUNetPlans__3d_fullres", "dataset_json_fallback": f"{RAW_ROOT}/Dataset003_BMLMPS_T1CE/dataset.json", }, "Dataset004_BrainMets": { "plans_folder": "nnUNetCLSTrainerMTL__nnUNetPlans__3d_fullres", "dataset_json_fallback": f"{RAW_ROOT}/Dataset004_BrainMets/dataset.json", }, "Dataset005_MU_Glioma_Post": { "plans_folder": "nnUNetCLSTrainerMTL__nnUNetPlans__3d_fullres", "dataset_json_fallback": f"{RAW_ROOT}/Dataset005_MU_Glioma_Post/dataset.json", }, "Dataset006_JSC_UCSD_PTGB": { "plans_folder": "nnUNetCLSTrainerMTL__nnUNetPlans__3d_fullres", "dataset_json_fallback": f"{RAW_ROOT}/Dataset006_JSC_UCSD_PTGB/dataset.json", }, "Dataset007_PICAI": { "plans_folder": "nnUNetCLSTrainerMTL__nnUNetPlans__3d_fullres", "dataset_json_fallback": f"{RAW_ROOT}/Dataset007_PICAI/dataset.json", }, "Dataset008_PETWB_Liver": { "plans_folder": "nnUNetCLSTrainerMTL__nnUNetPlans__3d_fullres", "dataset_json_fallback": f"{RAW_ROOT}/Dataset008_PETWB_Liver/dataset.json", }, "Dataset009_LUNA25": { "plans_folder": "nnUNetCLSTrainerMTL__nnUNetPlans__3d_fullres", "dataset_json_fallback": f"{RAW_ROOT}/Dataset009_LUNA25/dataset.json", }, } # ─── Single-fold predictor (adapted from fivefold_eval.py) ───────────────────── def _logit_to_segment(predicted_logits): max_logit, max_class = torch.max(predicted_logits, dim=0) return torch.where(max_logit >= 0.5, max_class, torch.tensor(0, device=predicted_logits.device)) def _convert_predicted_logits_to_segmentation_with_correct_shape( predicted_logits, plans_manager, configuration_manager, label_manager, properties_dict, use_softmax): spacing_transposed = [properties_dict['spacing'][i] for i in plans_manager.transpose_forward] current_spacing = configuration_manager.spacing if \ len(configuration_manager.spacing) == \ len(properties_dict['shape_after_cropping_and_before_resampling']) else \ [spacing_transposed[0], *configuration_manager.spacing] if properties_dict['shape_after_cropping_and_before_resampling'][0] < 600: predicted_logits = fast_resample_logit_to_shape( predicted_logits, properties_dict['shape_after_cropping_and_before_resampling'], current_spacing, [properties_dict['spacing'][i] for i in plans_manager.transpose_forward]) gc.collect() empty_cache(predicted_logits.device) if use_softmax: predicted_probabilities = label_manager.apply_inference_nonlin(predicted_logits) del predicted_logits segmentation = label_manager.convert_probabilities_to_segmentation(predicted_probabilities) else: segmentation = _logit_to_segment(predicted_logits) else: segmentation = fast_resample_logit_to_shape( predicted_logits, properties_dict['shape_after_cropping_and_before_resampling'], current_spacing, [properties_dict['spacing'][i] for i in plans_manager.transpose_forward]) dtype = torch.uint8 if len(label_manager.foreground_labels) < 255 else torch.uint16 seg_reverted = torch.zeros(properties_dict['shape_before_cropping'], dtype=dtype) slicer = bounding_box_to_slice(properties_dict['bbox_used_for_cropping']) seg_reverted[slicer] = segmentation del segmentation seg_reverted = seg_reverted.permute(plans_manager.transpose_backward) return seg_reverted.cpu() class SingleFoldPredictor(nnUNetPredictor): """Loads a single fold of an MTL (seg+cls) nnUNet model.""" def initialize_from_trained_model_folder(self, model_training_output_dir: str, fold: int, checkpoint_name: str): dataset_json = load_json(join(model_training_output_dir, 'dataset.json')) plans = load_json(join(model_training_output_dir, 'plans.json')) plans_manager = PlansManager(plans) ckpt_path = join(model_training_output_dir, f'fold_{fold}', checkpoint_name) checkpoint = torch.load(ckpt_path, map_location='cpu', weights_only=False) trainer_name = checkpoint['trainer_name'] configuration_name = checkpoint['init_args']['configuration'] inference_allowed_mirroring_axes = checkpoint.get('inference_allowed_mirroring_axes', None) weights = {k.replace('module.', ''): v for k, v in checkpoint['network_weights'].items()} safe_state_dict = {} for k, v in weights.items(): if any(x in k for x in ['running_mean', 'running_var', 'num_batches_tracked']): safe_state_dict[k] = v.clone() else: safe_state_dict[k] = v configuration_manager = plans_manager.get_configuration(configuration_name) if 'cls_class_num' in checkpoint: self.cls_class_num = checkpoint['cls_class_num'] else: cls_weight_keys = [k for k in safe_state_dict if 'classifier' in k and 'weight' in k] last_cls_key = sorted(cls_weight_keys)[-1] self.cls_class_num = safe_state_dict[last_cls_key].shape[0] num_input_channels = determine_num_input_channels(plans_manager, configuration_manager, dataset_json) trainer_class = recursive_find_python_class( join(nnunetv2.__path__[0], "training", "nnUNetTrainer"), trainer_name, 'nnunetv2.training.nnUNetTrainer') if trainer_class is None: raise RuntimeError(f'Unable to locate trainer class {trainer_name}') cls_head_output = self.cls_class_num if self.cls_class_num > 2 else 1 network = trainer_class.build_network_architecture( configuration_manager.network_arch_class_name, configuration_manager.network_arch_init_kwargs, configuration_manager.network_arch_init_kwargs_req_import, num_input_channels, plans_manager.get_label_manager(dataset_json).num_segmentation_heads, enable_deep_supervision=False, emb_dim=320, cls_class_num=cls_head_output, ) self.plans_manager = plans_manager self.configuration_manager = configuration_manager self.network = network self.network.load_state_dict(safe_state_dict) self.dataset_json = dataset_json self.trainer_name = trainer_name self.allowed_mirroring_axes = inference_allowed_mirroring_axes self.label_manager = plans_manager.get_label_manager(dataset_json) def preprocess(self, image, props): preprocessor = self.configuration_manager.preprocessor_class(verbose=False) data = preprocessor.run_case_npy(image, None, props, self.plans_manager, self.configuration_manager, self.dataset_json) return torch.from_numpy(data[0]).to(dtype=torch.float32, memory_format=torch.contiguous_format) @torch.inference_mode() def _internal_maybe_mirror_and_predict(self, x: torch.Tensor): mirror_axes = self.allowed_mirroring_axes if self.use_mirroring else None prediction, cls_prediction = self.network(x) if mirror_axes is not None: assert max(mirror_axes) <= x.ndim - 3 mirror_axes = [m + 2 for m in mirror_axes] axes_combinations = [ c for i in range(len(mirror_axes)) for c in itertools.combinations(mirror_axes, i + 1) ] for axes in axes_combinations: pred_f, cls_f = self.network(torch.flip(x, axes)) prediction += torch.flip(pred_f, axes) cls_prediction += cls_f prediction /= (len(axes_combinations) + 1) cls_prediction /= (len(axes_combinations) + 1) return prediction, cls_prediction @torch.inference_mode() def _internal_predict_sliding_window_return_logits(self, data, slicers, do_on_device=True): results_device = self.device if do_on_device else torch.device('cpu') self.network = self.network.to(self.device) self.network.eval() empty_cache(self.device) data = data.to(results_device) predicted_logits = torch.zeros( (self.label_manager.num_segmentation_heads, *data.shape[1:]), dtype=torch.half, device=results_device) n_predictions = torch.zeros(data.shape[1:], dtype=torch.half, device=results_device) cls_out_dim = self.cls_class_num if self.cls_class_num > 2 else 1 class_logits = torch.zeros((cls_out_dim,), dtype=torch.half, device=results_device) if self.use_gaussian: gaussian = compute_gaussian(tuple(self.configuration_manager.patch_size), sigma_scale=1. / 8, value_scaling_factor=10, device=results_device) else: gaussian = 1 for sl in tqdm(slicers, disable=not self.allow_tqdm, desc="sliding window"): workon = data[sl][None].to(self.device) pred, cls_patch = self._internal_maybe_mirror_and_predict(workon) pred = pred[0].to(results_device) cls_patch = cls_patch[0].to(results_device) if self.use_gaussian: pred *= gaussian predicted_logits[sl] += pred n_predictions[sl[1:]] += gaussian class_logits += cls_patch predicted_logits /= n_predictions class_logits /= len(slicers) if torch.any(torch.isinf(predicted_logits)): raise RuntimeError('Encountered inf in predicted array.') return predicted_logits, class_logits @torch.inference_mode() def inference(self, image, properties_dict, use_softmax: bool = False): image = self.preprocess(image, properties_dict) empty_cache(self.device) data, slicer_revert_padding = pad_nd_image( image, self.configuration_manager.patch_size, 'constant', {'value': 0}, True, None) slicers = self._internal_get_sliding_window_slicers(data.shape[1:]) seg_logit, class_logit = self._internal_predict_sliding_window_return_logits( data, slicers, self.perform_everything_on_device) if self.cls_class_num > 2: cls_probs = torch.softmax(class_logit.float(), dim=0).cpu() else: cls_probs = torch.sigmoid(class_logit.float()).cpu() seg_logit = seg_logit.cpu() empty_cache(self.device) seg_logit = seg_logit[(slice(None), *slicer_revert_padding[1:])] segmentation = _convert_predicted_logits_to_segmentation_with_correct_shape( seg_logit, self.plans_manager, self.configuration_manager, self.label_manager, properties_dict, use_softmax) return segmentation, cls_probs # ─── Classification label lookup ─────────────────────────────────────────────── def _load_dataset_json(model_dir: str, fallback_path: str) -> dict: """Return dataset.json, preferring the one packaged with the model.""" primary = join(model_dir, 'dataset.json') data = load_json(primary) if 'classification_labels' not in data and os.path.isfile(fallback_path): fallback = load_json(fallback_path) if 'classification_labels' in fallback: data['classification_labels'] = fallback['classification_labels'] return data def _format_cls_results(cls_probs: torch.Tensor, dataset_json: dict) -> dict: """Build {task_name: {class_name: probability}} from raw cls probs.""" cls_labels = dataset_json.get('classification_labels', {}) if not cls_labels: return {"unknown": {str(i): float(p) for i, p in enumerate(cls_probs.flatten().tolist())}} task_name = list(cls_labels.keys())[0] name_map = cls_labels[task_name] probs = cls_probs.flatten().tolist() if len(probs) == 1: p_pos = float(probs[0]) return {task_name: {name_map["0"]: 1.0 - p_pos, name_map["1"]: p_pos}} # Multi-class. Model may output more classes than dataset.json names # (e.g. Dataset004/006 trained with an extra "Unknown" bucket for -1 labels). out = {} for i, p in enumerate(probs): name = name_map.get(str(i), "Unknown" if i == len(name_map) else f"class_{i}") out[name] = float(p) return {task_name: out} # ─── Video overlay helper ─────────────────────────────────────────────────────── # Distinct BGR colors (OpenCV uses BGR) _LABEL_COLORS = [ (0, 0, 255), # red (0, 255, 0), # green (255, 0, 0), # blue (0, 255, 255), # yellow (255, 0, 255), # magenta (255, 255, 0), # cyan ] def _normalize_slice_uint8(img2d: np.ndarray) -> np.ndarray: img = img2d.astype(np.float32) lo, hi = np.percentile(img, (1, 99)) if hi <= lo: lo, hi = float(img.min()), float(img.max() if img.max() > img.min() else img.min() + 1) img = np.clip(img, lo, hi) img = (img - lo) / (hi - lo) * 255.0 return img.astype(np.uint8) def _overlay_slice(gray_u8: np.ndarray, seg2d: np.ndarray, label_values: Sequence[int], alpha: float = 0.4) -> np.ndarray: rgb = cv2.cvtColor(gray_u8, cv2.COLOR_GRAY2BGR) overlay = rgb.copy() for i, lbl in enumerate(label_values): mask = seg2d == lbl if mask.any(): overlay[mask] = _LABEL_COLORS[i % len(_LABEL_COLORS)] return cv2.addWeighted(overlay, alpha, rgb, 1 - alpha, 0) def _build_frames(image_zyx: np.ndarray, seg_zyx: np.ndarray, label_values: Sequence[int], label_names: Sequence[str]): """Yield (H, W, 3) BGR uint8 frames: axial slice + seg overlay + legend.""" assert image_zyx.shape == seg_zyx.shape, \ f"image/seg shape mismatch: {image_zyx.shape} vs {seg_zyx.shape}" z, h, w = image_zyx.shape legend_h = 22 + 18 * ((len(label_values) + 2) // 3) frame_h = h + legend_h # ffmpeg's libx264 requires even dimensions. frame_h += frame_h % 2 frame_w = w + w % 2 font = cv2.FONT_HERSHEY_SIMPLEX for zi in range(z): gray = _normalize_slice_uint8(image_zyx[zi]) frame_rgb = _overlay_slice(gray, seg_zyx[zi], label_values) canvas = np.zeros((frame_h, frame_w, 3), dtype=np.uint8) canvas[:h, :w] = frame_rgb cv2.putText(canvas, f"slice {zi + 1}/{z}", (6, 16), font, 0.5, (255, 255, 255), 1, cv2.LINE_AA) for i, (lbl, name) in enumerate(zip(label_values, label_names)): row, col = divmod(i, 3) x0 = 6 + col * (frame_w // 3) y0 = h + 16 + row * 18 color = _LABEL_COLORS[i % len(_LABEL_COLORS)] cv2.rectangle(canvas, (x0, y0 - 10), (x0 + 14, y0 + 2), color, -1) cv2.putText(canvas, name, (x0 + 18, y0), font, 0.4, (255, 255, 255), 1, cv2.LINE_AA) yield canvas def _write_overlay_video(image_zyx: np.ndarray, seg_zyx: np.ndarray, label_values: Sequence[int], label_names: Sequence[str], out_path: str, fps: int = 10) -> str: """Write a browser-playable H.264 mp4 by piping raw BGR frames to ffmpeg. Falls back to OpenCV's mp4v VideoWriter if ffmpeg is unavailable. """ os.makedirs(os.path.dirname(os.path.abspath(out_path)) or '.', exist_ok=True) frames = list(_build_frames(image_zyx, seg_zyx, label_values, label_names)) if not frames: raise RuntimeError("No frames to write") fh, fw = frames[0].shape[:2] ffmpeg_bin = shutil.which("ffmpeg") if ffmpeg_bin is not None: cmd = [ ffmpeg_bin, "-y", "-f", "rawvideo", "-vcodec", "rawvideo", "-s", f"{fw}x{fh}", "-pix_fmt", "bgr24", "-r", str(fps), "-i", "-", "-c:v", "libx264", "-pix_fmt", "yuv420p", "-preset", "medium", "-crf", "23", "-movflags", "+faststart", out_path, ] proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE) try: for frame in frames: proc.stdin.write(frame.tobytes()) proc.stdin.close() rc = proc.wait() if rc != 0: err = proc.stderr.read().decode('utf-8', errors='replace') raise RuntimeError(f"ffmpeg exited with {rc}:\n{err}") finally: if proc.stderr: proc.stderr.close() return out_path # Fallback: OpenCV mp4v (may not play in all viewers) fourcc = cv2.VideoWriter_fourcc(*'mp4v') writer = cv2.VideoWriter(out_path, fourcc, fps, (fw, fh)) if not writer.isOpened(): raise RuntimeError(f"Could not open video writer for {out_path}") for frame in frames: writer.write(frame) writer.release() return out_path # ─── Core inference runner ────────────────────────────────────────────────────── _PREDICTOR_CACHE: dict = {} def _get_predictor(dataset_name: str, device: Union[str, torch.device] = 'cuda', fold: int = 0, checkpoint: str = 'checkpoint_best.pth') -> SingleFoldPredictor: cache_key = (dataset_name, str(device), fold, checkpoint) if cache_key in _PREDICTOR_CACHE: return _PREDICTOR_CACHE[cache_key] reg = MODEL_REGISTRY[dataset_name] model_dir = join(BASELINE_ROOT, dataset_name, reg['plans_folder']) dev = torch.device(device, 0) if isinstance(device, str) and device != 'cpu' else \ (torch.device('cpu') if device == 'cpu' else device) perform_on_device = (dev.type != 'cpu') predictor = SingleFoldPredictor( tile_step_size=0.5, use_gaussian=True, use_mirroring=True, perform_everything_on_device=perform_on_device, device=dev, verbose=False, verbose_preprocessing=False, allow_tqdm=True, ) predictor.initialize_from_trained_model_folder(model_dir, fold=fold, checkpoint_name=checkpoint) predictor.network.to(dev) # Prefer the raw-dataset dataset.json when classification_labels is missing predictor.dataset_json = _load_dataset_json(model_dir, reg['dataset_json_fallback']) _PREDICTOR_CACHE[cache_key] = predictor return predictor def _as_path_list(image: Union[str, Sequence[str]]) -> List[str]: if isinstance(image, str): return [image] return list(image) def _run_inference(dataset_name: str, image: Union[str, Sequence[str]], output_dir: str, device: Union[str, torch.device] = 'cuda', fold: int = 0, use_softmax: bool = False, case_id: Optional[str] = None) -> Tuple[str, str, dict]: """Shared body used by every per-dataset function below.""" os.makedirs(output_dir, exist_ok=True) image_paths = _as_path_list(image) predictor = _get_predictor(dataset_name, device=device, fold=fold) expected = len(predictor.dataset_json['channel_names']) if len(image_paths) != expected: raise ValueError( f"{dataset_name} expects {expected} modality file(s) " f"(channels={list(predictor.dataset_json['channel_names'].values())}), " f"but got {len(image_paths)}: {image_paths}" ) if case_id is None: first = os.path.basename(image_paths[0]) for suffix in ('.nii.gz', '.nii', '.mha', '.nrrd'): if first.endswith(suffix): first = first[: -len(suffix)] break # strip trailing channel index like _0000 if len(first) > 5 and first[-5] == '_' and first[-4:].isdigit(): first = first[:-5] case_id = first image_npy, props = SimpleITKIO().read_images(image_paths) segmentation, cls_probs = predictor.inference(image_npy, props, use_softmax=use_softmax) seg_np = segmentation.numpy() # ─── Save segmentation NIfTI ────────────────────────────────────────── seg_path = join(output_dir, f"{case_id}_seg.nii.gz") sitk_img = sitk.GetImageFromArray(seg_np.astype(np.uint8)) sitk_img.SetSpacing(props['sitk_stuff']['spacing']) sitk_img.SetOrigin(props['sitk_stuff']['origin']) sitk_img.SetDirection(props['sitk_stuff']['direction']) sitk.WriteImage(sitk_img, seg_path) # ─── Build overlay video from the FIRST modality ───────────────────── label_map = predictor.dataset_json['labels'] # name -> value fg = [(name, val) for name, val in label_map.items() if val != 0] fg.sort(key=lambda x: x[1]) label_values = [v for _, v in fg] label_names = [n for n, _ in fg] first_modality_zyx = image_npy[0] # read_images returns (C, Z, Y, X) video_path = join(output_dir, f"{case_id}_overlay.mp4") _write_overlay_video(first_modality_zyx, seg_np, label_values, label_names, video_path) # ─── Classification results ────────────────────────────────────────── cls_results = _format_cls_results(cls_probs, predictor.dataset_json) cls_path = join(output_dir, f"{case_id}_classification.json") with open(cls_path, 'w') as f: json.dump(cls_results, f, indent=2) return seg_path, video_path, cls_results # ─── Per-model public functions ───────────────────────────────────────────────── def infer_dataset001_petwb_lung(image: Union[str, Sequence[str]], output_dir: str, device: Union[str, torch.device] = 'cuda', fold: int = 0) -> Tuple[str, str, dict]: """PETWB Lung — lung seg + cancer type (Other / Lung Cancer).""" return _run_inference("Dataset001_PETWB_Lung", image, output_dir, device, fold) def infer_dataset002_bmlmps_flair(image: Union[str, Sequence[str]], output_dir: str, device: Union[str, torch.device] = 'cuda', fold: int = 0) -> Tuple[str, str, dict]: """BMLMPS FLAIR — whole-tumor seg + EGFR status (Wild-Type / Mutation).""" return _run_inference("Dataset002_BMLMPS_FLAIR", image, output_dir, device, fold) def infer_dataset003_bmlmps_t1ce(image: Union[str, Sequence[str]], output_dir: str, device: Union[str, torch.device] = 'cuda', fold: int = 0) -> Tuple[str, str, dict]: """BMLMPS T1CE — core-tumor seg + EGFR status (Wild-Type / Mutation).""" return _run_inference("Dataset003_BMLMPS_T1CE", image, output_dir, device, fold) def infer_dataset004_brainmets(image: Union[str, Sequence[str]], output_dir: str, device: Union[str, torch.device] = 'cuda', fold: int = 0) -> Tuple[str, str, dict]: """PROTEAS BrainMets — necrotic/enhancing/edema seg + primary tumor origin (NSCLC / Breast).""" return _run_inference("Dataset004_BrainMets", image, output_dir, device, fold) def infer_dataset005_mu_glioma_post(image: Union[str, Sequence[str]], output_dir: str, device: Union[str, torch.device] = 'cuda', fold: int = 0) -> Tuple[str, str, dict]: """MU-Glioma-Post — NCR/ED/ET/NET_RC seg + primary diagnosis (GBM / Astrocytoma / Others).""" return _run_inference("Dataset005_MU_Glioma_Post", image, output_dir, device, fold) def infer_dataset006_jsc_ucsd_ptgb(image: Union[str, Sequence[str]], output_dir: str, device: Union[str, torch.device] = 'cuda', fold: int = 0) -> Tuple[str, str, dict]: """UCSD Post-Tx GBM — tumor seg + IDH mutation status (Wild-Type / Mutant).""" return _run_inference("Dataset006_JSC_UCSD_PTGB", image, output_dir, device, fold) def infer_dataset007_picai(image: Union[str, Sequence[str]], output_dir: str, device: Union[str, torch.device] = 'cuda', fold: int = 0) -> Tuple[str, str, dict]: """PI-CAI — csPCa seg + ISUP grade (6-class).""" return _run_inference("Dataset007_PICAI", image, output_dir, device, fold) def infer_dataset008_petwb_liver(image: Union[str, Sequence[str]], output_dir: str, device: Union[str, torch.device] = 'cuda', fold: int = 0) -> Tuple[str, str, dict]: """PETWB Liver — liver seg + cancer type (Other / Liver Cancer).""" return _run_inference("Dataset008_PETWB_Liver", image, output_dir, device, fold) def infer_dataset009_luna25(image: Union[str, Sequence[str]], output_dir: str, device: Union[str, torch.device] = 'cuda', fold: int = 0) -> Tuple[str, str, dict]: """LUNA25 — pulmonary nodule seg + malignancy (Benign / Malignant).""" return _run_inference("Dataset009_LUNA25", image, output_dir, device, fold) DATASET_DISPATCH = { "Dataset001_PETWB_Lung": infer_dataset001_petwb_lung, "Dataset002_BMLMPS_FLAIR": infer_dataset002_bmlmps_flair, "Dataset003_BMLMPS_T1CE": infer_dataset003_bmlmps_t1ce, "Dataset004_BrainMets": infer_dataset004_brainmets, "Dataset005_MU_Glioma_Post": infer_dataset005_mu_glioma_post, "Dataset006_JSC_UCSD_PTGB": infer_dataset006_jsc_ucsd_ptgb, "Dataset007_PICAI": infer_dataset007_picai, "Dataset008_PETWB_Liver": infer_dataset008_petwb_liver, "Dataset009_LUNA25": infer_dataset009_luna25, } # ─── CLI ──────────────────────────────────────────────────────────────────────── if __name__ == "__main__": import argparse parser = argparse.ArgumentParser( description="Run fold-0 inference for a AutoMSC_challenge baseline model.") parser.add_argument('--dataset', required=True, choices=list(DATASET_DISPATCH.keys()), help='Baseline dataset to use.') parser.add_argument('--input', required=True, nargs='+', help='One path per modality, in the channel order from dataset.json.') parser.add_argument('--output_dir', required=True, help='Where to write seg/video/cls files.') parser.add_argument('--device', default='cuda') parser.add_argument('--fold', type=int, default=0) args = parser.parse_args() fn = DATASET_DISPATCH[args.dataset] seg_path, video_path, cls_results = fn( image=args.input, output_dir=args.output_dir, device=args.device, fold=args.fold, ) print(f"Segmentation: {seg_path}") print(f"Overlay video: {video_path}") print(f"Classification: {json.dumps(cls_results, indent=2)}")