| """ |
| Loading the released weights. |
| |
| Every file under weights/ is a plain dict of tensors + primitives, so it loads |
| under `torch.load(..., weights_only=True)` -- no custom classes, no pickled |
| training objects, no need to have this repo importable to read it. |
| |
| Layout of a checkpoint file:: |
| |
| { |
| "state_dict": OrderedDict[str, Tensor], # no "module." prefix |
| "arch": "c3d" | "mednext" | "swinunetr", |
| "in_channels": 8 (pretrained) or 5 (finetuned / fromscratch), |
| "out_channels": 1, |
| "anatomy": "han" | "pancreas" | "source-domain", |
| "regime": "pretrained" | "finetuned" | "fromscratch", |
| "act_sig": True, # apply sigmoid to the raw output |
| "scale_out": 7.5 | 5.5, # then multiply by this |
| "dose_div_factor": 10, # then multiply by this -> dose in Gy |
| "source": "...", "note": "...", |
| } |
| |
| So the raw network output becomes physical dose as:: |
| |
| dose_Gy = sigmoid(out) * scale_out * dose_div_factor |
| """ |
|
|
| import torch |
|
|
| from build_model import build_model |
|
|
|
|
| def load_checkpoint(path, map_location="cpu"): |
| """Read a released checkpoint file. Returns the full dict.""" |
| return torch.load(path, map_location=map_location, weights_only=True) |
|
|
|
|
| def load_model(path, device="cpu", strict=True): |
| """ |
| Build the right architecture for a checkpoint and load its weights. |
| |
| Returns ``(model_in_eval_mode, checkpoint_dict)``. The checkpoint dict |
| carries the post-processing constants (``act_sig``, ``scale_out``, |
| ``dose_div_factor``) you need to turn the output into Gy. |
| """ |
| ckpt = load_checkpoint(path, map_location="cpu") |
|
|
| model = build_model( |
| ckpt["arch"], |
| in_channels=ckpt["in_channels"], |
| out_channels=ckpt.get("out_channels", 1), |
| ) |
| model.load_state_dict(ckpt["state_dict"], strict=strict) |
| model.to(device).eval() |
| return model, ckpt |
|
|
|
|
| def to_dose(output, ckpt): |
| """ |
| Convert a raw network output to dose in Gy using the checkpoint's own |
| constants. Handles C3D's ``[output_A, output_B]`` and MedNeXt's optional |
| deep-supervision list by taking the refined / main head. |
| """ |
| if isinstance(output, (list, tuple)): |
| |
| output = output[-1] if ckpt["arch"] == "c3d" else output[0] |
|
|
| if ckpt.get("act_sig", True): |
| output = torch.sigmoid(output) |
|
|
| return output * ckpt["scale_out"] * ckpt["dose_div_factor"] |
|
|
