| import argparse |
| import csv |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| IMAGENET_MEAN = (0.485, 0.456, 0.406) |
| IMAGENET_STD = (0.229, 0.224, 0.225) |
| CLAHE_CLIP, CLAHE_GRID = 1.0, (8, 8) |
| DEFAULT_WEIGHTS = Path(__file__).resolve().parent / "weights" |
| DEFAULT_CONFIG = Path(__file__).resolve().parent / "config.json" |
|
|
|
|
| def load_config(config_path=DEFAULT_CONFIG): |
| return json.loads(Path(config_path).read_text()) |
|
|
|
|
| def minmax_uint8(arr): |
| arr = np.asarray(arr).astype(np.float32) |
| mn, mx = float(arr.min()), float(arr.max()) |
| return ((arr - mn) / (mx - mn + 1e-8) * 255.0).astype(np.uint8) |
|
|
|
|
| def _as_dataset(dcm): |
| import pydicom |
|
|
| if isinstance(dcm, (str, bytes)) or hasattr(dcm, "__fspath__"): |
| return pydicom.dcmread(dcm, force=True) |
| return dcm |
|
|
|
|
| def dicom_to_image(dcm): |
| import cv2 |
|
|
| ds = _as_dataset(dcm) |
| raw = ds.pixel_array.astype(np.float32) |
| if raw.ndim != 2: |
| raise ValueError(f"Expected a single grayscale X-ray, got pixel shape {raw.shape}") |
| if str(getattr(ds, "PhotometricInterpretation", "")) == "MONOCHROME1": |
| raw = raw.max() - raw |
| base = minmax_uint8(raw) |
| eq = cv2.equalizeHist(base) |
| clahe = cv2.createCLAHE(clipLimit=CLAHE_CLIP, tileGridSize=CLAHE_GRID).apply(base) |
| return np.stack([eq, clahe, base], axis=-1), base.shape |
|
|
|
|
| def letterbox_topleft(img, size): |
| import cv2 |
|
|
| h, w = img.shape[:2] |
| scale = size / max(h, w) |
| nh, nw = int(round(h * scale)), int(round(w * scale)) |
| resized = cv2.resize(img, (nw, nh), interpolation=cv2.INTER_LINEAR) |
| padded = np.zeros((size, size, img.shape[2]), dtype=img.dtype) |
| padded[:nh, :nw, :] = resized |
| return padded, (nh, nw) |
|
|
|
|
| def pad_square_topleft(img, fill=0): |
| h, w = img.shape[:2] |
| if h == w: |
| return img |
| s = max(h, w) |
| out = np.full((s, s, img.shape[2]), fill, dtype=img.dtype) |
| out[:h, :w, :] = img |
| return out |
|
|
|
|
| def to_input_resized(img3, size, mean, std, device, dtype): |
| x = torch.from_numpy(np.ascontiguousarray(img3)).permute(2, 0, 1)[None].float() |
| if tuple(x.shape[-2:]) != (size, size): |
| x = F.interpolate(x, size=(size, size), mode="bilinear", align_corners=False, antialias=True) |
| x = x / 255.0 |
| x = (x - torch.tensor(mean).view(1, 3, 1, 1)) / torch.tensor(std).view(1, 3, 1, 1) |
| return x.to(device=device, dtype=dtype) |
|
|
|
|
| def to_input(img3, mean, std, device, dtype): |
| x = torch.from_numpy(np.ascontiguousarray(img3)).float().permute(2, 0, 1) / 255.0 |
| x = (x - torch.tensor(mean).view(3, 1, 1)) / torch.tensor(std).view(3, 1, 1) |
| return x.unsqueeze(0).to(device=device, dtype=dtype) |
|
|
|
|
| def present_fallback_mask(prob, fraction): |
| prob = np.asarray(prob, dtype=np.float32) |
| if float(prob.max()) <= 0.0: |
| return np.zeros(prob.shape, dtype=np.uint8) |
| n = max(1, int(prob.size * fraction)) |
| mask = np.zeros(prob.size, np.uint8) |
| mask[np.argpartition(prob.ravel(), -n)[-n:]] = 1 |
| return mask.reshape(prob.shape) |
|
|
|
|
| def _group_norm(num_channels): |
| groups = next(g for g in (32, 16, 8, 4, 2, 1) if num_channels % g == 0) |
| return nn.GroupNorm(groups, num_channels) |
|
|
|
|
| class ConvNormAct(nn.Module): |
| def __init__(self, in_ch, out_ch, k=3): |
| super().__init__() |
| self.block = nn.Sequential(nn.Conv2d(in_ch, out_ch, k, padding=k // 2, bias=False), _group_norm(out_ch), nn.GELU()) |
|
|
| def forward(self, x): |
| return self.block(x) |
|
|
|
|
| class DoubleConv(nn.Module): |
| def __init__(self, in_ch, out_ch): |
| super().__init__() |
| self.block = nn.Sequential(ConvNormAct(in_ch, out_ch), ConvNormAct(out_ch, out_ch)) |
|
|
| def forward(self, x): |
| return self.block(x) |
|
|
|
|
| class PPM(nn.Module): |
| def __init__(self, in_channels, out_channels, pool_sizes=(1, 2, 3, 6), reduction=4): |
| super().__init__() |
| inter = max(in_channels // reduction, 32) |
| self.stages = nn.ModuleList( |
| [nn.Sequential(nn.AdaptiveAvgPool2d(ps), ConvNormAct(in_channels, inter, k=1)) for ps in pool_sizes] |
| ) |
| self.bottleneck = ConvNormAct(in_channels + inter * len(pool_sizes), out_channels, k=3) |
|
|
| def forward(self, x): |
| h, w = x.shape[-2:] |
| feats = [x] |
| for stage in self.stages: |
| feats.append(F.interpolate(stage(x), size=(h, w), mode="bilinear", align_corners=False)) |
| return self.bottleneck(torch.cat(feats, dim=1)) |
|
|
|
|
| class ConvNextUNetDecoder(nn.Module): |
| def __init__(self, encoder_channels, decoder_channels, out_channels, target_mask_size, ppm_pool_sizes): |
| super().__init__() |
| self.target_mask_size = target_mask_size |
| self.bottleneck = PPM(encoder_channels[0], decoder_channels[0], pool_sizes=ppm_pool_sizes) |
| self.up_blocks = nn.ModuleList([ |
| DoubleConv(decoder_channels[i] + encoder_channels[i + 1], decoder_channels[i + 1]) |
| for i in range(len(encoder_channels) - 1) |
| ]) |
| self.final_refine = DoubleConv(decoder_channels[-1], decoder_channels[-1]) |
| self.head = nn.Conv2d(decoder_channels[-1], out_channels, 1) |
|
|
| def forward(self, feature_maps): |
| x = self.bottleneck(feature_maps[0]) |
| for up, skip in zip(self.up_blocks, feature_maps[1:]): |
| x = F.interpolate(x, size=skip.shape[-2:], mode="bilinear", align_corners=False) |
| x = up(torch.cat([x, skip], dim=1)) |
| x = self.final_refine(x) |
| x = F.interpolate(x, size=(self.target_mask_size, self.target_mask_size), mode="bilinear", align_corners=False) |
| return self.head(x) |
|
|
|
|
| class ConvNextTower(nn.Module): |
| def __init__(self, tower_cfg): |
| super().__init__() |
| from transformers import AutoConfig, AutoModel |
|
|
| cfg = AutoConfig.for_model(tower_cfg["model_type"], **{k: v for k, v in tower_cfg.items() if k != "model_type"}) |
| cfg.output_hidden_states = True |
| self.backbone = AutoModel.from_config(cfg) |
| hs = tuple(getattr(cfg, "hidden_sizes", None) or (192, 384, 768, 1536)) |
| self.encoder_channels = tuple(reversed(hs[-4:])) |
|
|
| def forward(self, pixel_values): |
| maps = list(self.backbone(pixel_values).hidden_states[-4:]) |
| return list(reversed(maps)) |
|
|
|
|
| def build_vit_tower(tower_cfg, patch_start): |
| from transformers import AutoConfig, AutoModel |
|
|
| expected = 1 + int(tower_cfg.get("num_register_tokens", 0)) |
| if int(patch_start) != expected: |
| raise ValueError(f"patch_start {patch_start} != 1 CLS + {expected - 1} register tokens") |
| cfg = AutoConfig.for_model(tower_cfg["model_type"], **{k: v for k, v in tower_cfg.items() if k != "model_type"}) |
| return AutoModel.from_config(cfg) |
|
|
|
|
| class SegModel(nn.Module): |
| def __init__(self, spec): |
| super().__init__() |
| self.vision_model = ConvNextTower(spec["tower"]) |
| self.segmentation_decoder = ConvNextUNetDecoder(self.vision_model.encoder_channels, **spec["decoder"]) |
|
|
| def forward(self, pixel_values): |
| return self.segmentation_decoder(self.vision_model(pixel_values)) |
|
|
|
|
| class TokenPooler(nn.Module): |
| def __init__(self, dim, num_heads): |
| super().__init__() |
| self.query = nn.Parameter(torch.zeros(1, 1, dim)) |
| self.attn = nn.MultiheadAttention(dim, num_heads, batch_first=True) |
| self.norm = nn.LayerNorm(dim) |
|
|
| def forward(self, tokens): |
| q = self.query.expand(tokens.shape[0], -1, -1) |
| pooled, _ = self.attn(q, tokens, tokens, need_weights=False) |
| return self.norm(pooled.squeeze(1)) |
|
|
|
|
| class ClassifierHead(nn.Module): |
| def __init__(self, in_channels, hidden_dim, num_classes, dropout): |
| super().__init__() |
| self.mlp = nn.Sequential( |
| nn.Linear(in_channels, hidden_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden_dim, num_classes) |
| ) |
|
|
| def forward(self, x): |
| return self.mlp(x).squeeze(-1) |
|
|
|
|
| class ClsModel(nn.Module): |
| def __init__(self, spec): |
| super().__init__() |
| self.patch_start = int(spec["patch_start"]) |
| self.vision_model = build_vit_tower(spec["tower"], self.patch_start) |
| self.pool = TokenPooler(spec["dim"], **spec["pool"]) |
| self.head = ClassifierHead(spec["dim"], **spec["head"]) |
|
|
| def forward(self, pixel_values): |
| tokens = self.vision_model(pixel_values).last_hidden_state[:, self.patch_start:, :] |
| return self.head(self.pool(tokens)) |
|
|
|
|
| class Task1Model: |
| def __init__(self, cls_net, seg_nets, seg_spec, cls_spec): |
| self.cls_net = cls_net |
| self.seg_nets = seg_nets |
| self.seg_size = int(seg_spec["infer_size"]) |
| self.cls_size = int(cls_spec["infer_size"]) |
| self.op = seg_spec["op"] |
| param = next(cls_net.parameters()) |
| self.device, self.dtype = param.device, param.dtype |
|
|
| @torch.no_grad() |
| def _cls_prob(self, dcm): |
| img3, _ = dicom_to_image(dcm) |
| x = to_input_resized(pad_square_topleft(img3), self.cls_size, IMAGENET_MEAN, IMAGENET_STD, |
| self.device, self.dtype) |
| return float(torch.sigmoid(self.cls_net(x).reshape(-1)[0].float())) |
|
|
| @torch.no_grad() |
| def _prob(self, dcm): |
| img3, native_hw = dicom_to_image(dcm) |
| padded, (nh, nw) = letterbox_topleft(img3, self.seg_size) |
| x = to_input(padded, IMAGENET_MEAN, IMAGENET_STD, self.device, self.dtype) |
| member_probs = [net(x).sigmoid()[0, 0] for net in self.seg_nets] |
| prob = torch.stack(member_probs).mean(0)[:nh, :nw] |
| prob = F.interpolate(prob[None, None].float(), size=native_hw, mode="bilinear", align_corners=False)[0, 0] |
| return prob.float().cpu().numpy() |
|
|
| @torch.no_grad() |
| def predict(self, dcm_path): |
| dcm = _as_dataset(dcm_path) |
| present = int(self._cls_prob(dcm) >= self.op["detection_tau"]) |
| if not self.seg_nets: |
| return None, present |
| prob = self._prob(dcm) |
| mask = (prob > self.op["threshold"]).astype(np.uint8) if present else np.zeros(prob.shape, np.uint8) |
| if present and not mask.any(): |
| mask = present_fallback_mask(prob, self.op["fallback_fraction"]) |
| return mask, present |
|
|
|
|
| def load_model(weights_path=DEFAULT_WEIGHTS, mode="seg", device="auto", config_path=DEFAULT_CONFIG): |
| from safetensors.torch import load_file |
|
|
| if mode not in ("seg", "cls"): |
| raise ValueError("mode must be 'seg' or 'cls'") |
| weights_path = Path(weights_path) |
| paths = [weights_path / f"model_{i}.safetensors" for i in range(5 if mode == "seg" else 1)] |
| for path in paths: |
| if not path.is_file(): |
| raise FileNotFoundError(f"Missing model weights: {path}") |
| config = load_config(config_path) |
| dev = torch.device("cuda" if torch.cuda.is_available() else "cpu") if device == "auto" else torch.device(device) |
| nets = [] |
| for i, path in enumerate(paths): |
| net = ClsModel(config["cls_spec"]) if i == 0 else SegModel(config["seg_spec"]) |
| net.load_state_dict(load_file(str(path)), strict=True) |
| nets.append(net.to(dev).eval()) |
| print(f"[weights] loaded {path.name}", flush=True) |
| return Task1Model(nets[0], nets[1:], config["seg_spec"], config["cls_spec"]) |
|
|
|
|
| def dicom_files(directory): |
| return sorted(p for p in directory.iterdir() if p.is_file() and p.suffix.lower() in (".dcm", ".dicom")) |
|
|
|
|
| def discover_cases(input_path): |
| path = Path(input_path) |
| if path.is_file(): |
| if path.suffix.lower() not in (".dcm", ".dicom"): |
| raise ValueError(f"Expected a .dcm or .dicom file: {path}") |
| return [(path.stem, path)] |
| if not path.is_dir(): |
| raise FileNotFoundError(f"Input does not exist: {path}") |
| files = dicom_files(path) |
| case_dirs = sorted(p for p in path.iterdir() if p.is_dir() and not p.name.startswith(".")) |
| if files and case_dirs: |
| raise ValueError("Input must contain either DICOM files or case folders, not both") |
| cases = [(p.stem, p) for p in files] |
| for directory in case_dirs: |
| paths = dicom_files(directory) |
| if not paths: |
| raise ValueError(f"No DICOM files found in case folder: {directory}") |
| cases.append((directory.name, paths[0])) |
| if not cases: |
| raise ValueError(f"No DICOM inputs found in: {path}") |
| ids = [case_id for case_id, _ in cases] |
| if len(set(ids)) != len(ids): |
| raise ValueError("Input files have duplicate case IDs") |
| return cases |
|
|
|
|
| def make_mask_match_reference(mask, reference): |
| mask = np.asarray(mask) |
| mask = (mask > 0).astype(np.uint8) |
|
|
| ref_dim = reference.GetDimension() |
| ref_size = reference.GetSize() |
|
|
| if ref_dim == 2: |
| expected_shape = (ref_size[1], ref_size[0]) |
| if mask.ndim == 3 and mask.shape[0] == 1: |
| mask = mask[0] |
| elif mask.ndim == 3 and mask.shape[-1] == 1: |
| mask = mask[..., 0] |
| if mask.ndim != 2: |
| raise RuntimeError(f"Expected 2D mask for 2D X-ray reference, got shape {mask.shape}") |
| if mask.shape != expected_shape: |
| raise RuntimeError(f"Mask shape {mask.shape} does not match reference shape {expected_shape}") |
| elif ref_dim == 3: |
| expected_shape = (ref_size[2], ref_size[1], ref_size[0]) |
| if mask.ndim == 2 and ref_size[2] == 1: |
| mask = mask[None, :, :] |
| if mask.ndim != 3: |
| raise RuntimeError(f"Expected 3D mask for 3D reference, got shape {mask.shape}") |
| if mask.shape != expected_shape: |
| raise RuntimeError(f"Mask shape {mask.shape} does not match reference shape {expected_shape}") |
| else: |
| raise RuntimeError(f"Unsupported reference dimension: {ref_dim}") |
|
|
| return mask.astype(np.uint8) |
|
|
|
|
| def save_mask(mask, dcm_path, output_path): |
| import SimpleITK as sitk |
|
|
| reference = sitk.ReadImage(str(dcm_path)) |
| mask = make_mask_match_reference(mask, reference) |
| pred = sitk.GetImageFromArray(mask) |
| pred.CopyInformation(reference) |
| sitk.WriteImage(pred, str(output_path)) |
|
|
|
|
| def run_one_case(model, case_id, dcm_path, output_dir): |
| mask, cavity = model.predict(dcm_path) |
| if mask is not None: |
| save_mask(mask, dcm_path, Path(output_dir) / f"{case_id}.nii.gz") |
| return case_id, cavity |
|
|
|
|
| def predict(input_path, output_dir, weights_path=DEFAULT_WEIGHTS, mode="seg", device="auto", config_path=DEFAULT_CONFIG): |
| cases = discover_cases(input_path) |
| model = load_model(weights_path, mode=mode, device=device, config_path=config_path) |
| output_dir = Path(output_dir) |
| output_dir.mkdir(parents=True, exist_ok=True) |
| rows = [run_one_case(model, case_id, path, output_dir) for case_id, path in cases] |
| with (output_dir / "prediction.csv").open("w", newline="") as f: |
| writer = csv.writer(f) |
| writer.writerow(["our_id", "cavity"]) |
| writer.writerows(rows) |
| print(f"[done] wrote {len(rows)} predictions (mode={mode}) -> {output_dir}", flush=True) |
| return rows |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="MultimodalAI: cavity detection and segmentation") |
| parser.add_argument("--input", default="/input", help="DICOM file, flat DICOM folder, or root of case folders") |
| parser.add_argument("--output", default="/output", help="output directory") |
| parser.add_argument("--weights", default=str(DEFAULT_WEIGHTS), help="directory containing model_0.safetensors through model_4.safetensors") |
| parser.add_argument("--config", default=str(DEFAULT_CONFIG), help="config.json holding model and threshold settings") |
| parser.add_argument("--mode", choices=("seg", "cls"), default="seg", help="seg: classification and masks (default); cls: classification only") |
| parser.add_argument("--device", default="auto", help="auto, cpu, cuda, or cuda:N") |
| args = parser.parse_args() |
| predict(args.input, args.output, args.weights, args.mode, args.device, args.config) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
