Model card / docstring: state the input scaling that actually works (clip [0,200] / 255)

README.md

@@ -1,193 +1,197 @@
----
-library_name: transformers
-license: mit
-tags:
-- vesuvius-challenge
-- ink-detection
-- herculaneum
-- resnet3d
-- u-net
-- 3d-segmentation
-- volumetric-imaging
-pipeline_tag: image-segmentation
----
-
-# PHerc.1667-iteration-4
-
-> **Trained on segment l_2 with l_2_inklabels4.png (24,773 tiles).**
-
-Ablation 4/5 — 24,773 training tiles.
-
-This is one of **six sibling models** released together — five label
-ablations on segment `l_2` (`ink1`–`ink5`, increasing label coverage)
-and one cross-segment baseline (`ink0`).  The full family is listed
-at the bottom of this card.
-
-## Preview
-
-`l_2` (training segment) prediction with the training label overlaid in
-magenta, and `l_5` (held-out segment) prediction. All panels are
-downsampled 16× and rotated 180° to match the publication-figure
-convention. The full-resolution `last.ckpt` outputs are at 43008 × ~30000
-voxels.
-
-| training label | l_2 prediction | l_5 prediction |
-|----------------|----------------|----------------|
-| ![label](./preview_label.png) | ![l_2 pred](./preview_l_2.png) | ![l_5 pred](./preview_l_5.png) |
-
-
-## Architecture in one paragraph
-
-A 3-D volumetric input `(B, 1, 62, 256, 256)` is encoded by a
-**ResNet3D-50** backbone (Hara, Kataoka & Satoh, 2018; initialised from
-the Kinetics-700 release `r3d50_KM_200ep.pth` with conv1 weights
-summed across RGB → 1 grayscale channel).  Each of the four backbone
-stages is collapsed along the z (depth) axis with `torch.max`,
-producing a 2-D feature pyramid `{(256,64,64), (512,32,32),
-(1024,16,16), (2048,8,8)}`.  A small **2-D U-Net decoder** upsamples
-coarse-to-fine with concatenated skip connections; a 1×1 conv head
-produces a single sigmoid logit channel at quarter resolution
-`(B, 1, 64, 64)`.  Training uses `0.5·Dice + 0.5·SoftBCE` against the
-label down-interpolated to 64×64.
-
-## Quick start
-
-```python
-import torch
-from transformers import AutoModel
-
-model = AutoModel.from_pretrained(
-    "YoussefMoNader/PHerc.1667-iteration-4",
-    trust_remote_code=True,
-).eval().cuda()
-
-# Input: float32, shape (B, 1, D=62, H=256, W=256).
-# Intensity should already be in roughly [0, 1] (the training pipeline
-# clipped raw uint8 layers to [0, 200] then applied Normalize(mean=0, std=1)
-# which keeps the magnitude small).
-x = torch.randn(1, 1, 62, 256, 256, device="cuda")
-
-with torch.no_grad():
-    out = model(x)
-
-print(out.logits.shape)                        # torch.Size([1, 1, 64, 64])
-prob = torch.sigmoid(out.logits)               # ink probability per pixel
-```
-
-## Full-segment inference (tiling)
-
-The model only sees 256×256 windows.  For a full scroll segment you
-need to slide the window across the (padded) layer stack and average
-overlapping predictions:
-
-```python
-import numpy as np, cv2, torch
-import torch.nn.functional as F
-from transformers import AutoModel
-
-model = AutoModel.from_pretrained(
-    "YoussefMoNader/PHerc.1667-iteration-4", trust_remote_code=True,
-).eval().cuda()
-
-WINDOW, STRIDE = 256, 128       # 128 = 2x oversample; 64 for 8x oversample
-D = 62                           # number of z-layers
-
-# image: (H, W, D) uint8 stack of the 62 layers, padded to multiples of 256.
-# fmask: (H, W) uint8 fragment mask (0 = outside, 255 = inside).
-H, W, _ = image.shape
-mask_pred  = np.zeros((H, W), dtype=np.float32)
-mask_count = np.zeros((H, W), dtype=np.float32)
-
-with torch.no_grad():
-    for y in range(0, H - WINDOW + 1, STRIDE):
-        for x in range(0, W - WINDOW + 1, STRIDE):
-            if np.any(fmask[y:y+WINDOW, x:x+WINDOW] == 0):
-                continue
-            tile = image[y:y+WINDOW, x:x+WINDOW]            # (256,256,62)
-            t = torch.from_numpy(tile).permute(2, 0, 1)     # (62,256,256)
-            t = t.unsqueeze(0).unsqueeze(0).float().cuda()  # (1,1,62,256,256)
-            logits = model(t).logits                        # (1,1,64,64)
-            prob   = torch.sigmoid(logits)
-            prob   = F.interpolate(prob, scale_factor=4,
-                                   mode="bilinear").squeeze().cpu().numpy()
-            mask_pred[y:y+WINDOW, x:x+WINDOW]  += prob
-            mask_count[y:y+WINDOW, x:x+WINDOW] += 1.0
-
-pred = np.divide(mask_pred, mask_count,
-                 out=np.zeros_like(mask_pred),
-                 where=mask_count != 0)
-cv2.imwrite("prediction.png", np.clip(pred * 255, 0, 255).astype(np.uint8))
-```
-
-## Training summary
-
-| | |
-|---|---|
-| **Backbone** | ResNet3D-50 (3-D conv, BN, ReLU residual blocks) |
-| **Encoder init** | `r3d50_KM_200ep.pth` (Kinetics-700), conv1 summed across RGB |
-| **Decoder** | 2-D U-Net (3 up-blocks: bilinear 2× + concat skip + 3×3 conv + BN + ReLU) |
-| **Output** | 1 channel, sigmoid logit, quarter-resolution (64×64) |
-| **Loss** | 0.5 × Dice + 0.5 × SoftBCE (smooth = 0.25) |
-| **Optimizer** | AdamW, OneCycle lr 2e-5 → 3e-4, pct_start = 0.15 |
-| **Batch** | 2 (effective 8 via accumulate 4), 16-mixed, grad-clip 1.0 |
-| **Max steps** | 12,396 (= 3 epochs over the densest ablation label) |
-| **Training segment(s)** | `l_2` |
-| **Training label** | `l_2_inklabels4.png` |
-| **Training tiles** (256×256 sub-tiles at stride 64) | **24,773** |
-| **Final train loss (`_epoch`)** | **0.5256** |
-| **Final train loss (`_step`, single-batch noise)** | 0.5005 |
-| **Wandb** | [vesuvius-challenge/paper/l2_ink4_l5infer](https://wandb.ai/vesuvius-challenge/paper/runs/h73rxb10) |
-| **Random seed** | 130697 |
-| **Determinism** | `cudnn.deterministic = True`, `cudnn.benchmark = False` |
-| **Hardware** | 1 × NVIDIA H100 80 GB; ≈ 2 h end-to-end (load + train + inference) |
-
-## Files
-
-| file | size | description |
-|------|------|-------------|
-| `config.json` | 1 KB | architecture + provenance metadata; loaded by `AutoConfig` |
-| `configuration_inkdetection.py` | 2 KB | `InkDetectionConfig(PretrainedConfig)` |
-| `modeling_inkdetection.py` | 9 KB | self-contained `InkDetectionModel(PreTrainedModel)` |
-| `model.safetensors` | 319 MB | converted weights (338 tensors) |
-| `last.ckpt` | 963 MB | original PyTorch-Lightning checkpoint (incl. optimizer + LR-scheduler state) — load with `torch.load(...)["state_dict"]` |
-| `preview_l_2.png` | ~700 KB | low-res preview of the l_2 prediction (1/16 scale, 180° rotated) |
-| `preview_l_5.png` | ~2 MB | low-res preview of the l_5 (held-out) prediction |
-| `preview_label.png` | ~50 KB | the training label, same scale + rotation |
-
-The HuggingFace weights are **bit-perfect identical** to the original
-PyTorch-Lightning checkpoint (verified `max abs diff = 0.0e+00` on
-identical inputs).  Use `model.safetensors` for `AutoModel.from_pretrained`;
-use `last.ckpt` only if you want to resume training from the saved
-optimizer / scheduler state.
-
-## The model family
-
-| model | training segment(s) | label | tiles | effective epochs |
-|-------|---------------------|-------|-------|-------------------|
-| [`PHerc.1667-iteration-0`](https://huggingface.co/scrollprize/PHerc.1667-iteration-0) | 500p2a + 658 + 20250910185200 + 20250919125754* | (cross-segment baseline) | 20,075 | ~5 |
-| [`PHerc.1667-iteration-1`](https://huggingface.co/scrollprize/PHerc.1667-iteration-1) | `l_2` | `l_2_inklabels.png`  | 3,396 | ~30 |
-| [`PHerc.1667-iteration-2`](https://huggingface.co/scrollprize/PHerc.1667-iteration-2) | `l_2` | `l_2_inklabels2.png` | 8,970 | ~12 |
-| [`PHerc.1667-iteration-3`](https://huggingface.co/scrollprize/PHerc.1667-iteration-3) | `l_2` | `l_2_inklabels3.png` | 15,286 | ~7 |
-| [`PHerc.1667-iteration-4`](https://huggingface.co/scrollprize/PHerc.1667-iteration-4) | `l_2` | `l_2_inklabels4.png` | 24,773 | ~5 |
-| [`PHerc.1667-iteration-5`](https://huggingface.co/scrollprize/PHerc.1667-iteration-5) | `l_2` | `l_2_inklabels5.png` | 33,061 | 3 |
-
-All six share the architecture, hyperparameters, and a fixed step
-budget of 12,396 optimizer steps; the only thing that varies between
-rows is the supervising label (or, for ink0, the training segments).
-
-## Citation
-
-If you use this model in published work, please cite the Vesuvius
-Challenge and the underlying ResNet3D paper:
-
-```bibtex
-@inproceedings{hara2018can,
-  title  = {Can spatiotemporal 3D CNNs retrace the history of 2D CNNs and ImageNet?},
-  author = {Hara, Kensho and Kataoka, Hirokatsu and Satoh, Yutaka},
-  booktitle = {CVPR}, year = {2018},
-}
-```
-
-## Licence
-
-MIT.
+---
+library_name: transformers
+license: mit
+tags:
+- vesuvius-challenge
+- ink-detection
+- herculaneum
+- resnet3d
+- u-net
+- 3d-segmentation
+- volumetric-imaging
+pipeline_tag: image-segmentation
+---
+
+# PHerc.1667-iteration-4
+
+> **Trained on segment l_2 with l_2_inklabels4.png (24,773 tiles).**
+
+Ablation 4/5 — 24,773 training tiles.
+
+This is one of **six sibling models** released together — five label
+ablations on segment `l_2` (`ink1`–`ink5`, increasing label coverage)
+and one cross-segment baseline (`ink0`).  The full family is listed
+at the bottom of this card.
+
+## Preview
+
+`l_2` (training segment) prediction with the training label overlaid in
+magenta, and `l_5` (held-out segment) prediction. All panels are
+downsampled 16× and rotated 180° to match the publication-figure
+convention. The full-resolution `last.ckpt` outputs are at 43008 × ~30000
+voxels.
+
+| training label | l_2 prediction | l_5 prediction |
+|----------------|----------------|----------------|
+| ![label](./preview_label.png) | ![l_2 pred](./preview_l_2.png) | ![l_5 pred](./preview_l_5.png) |
+
+
+## Architecture in one paragraph
+
+A 3-D volumetric input `(B, 1, 62, 256, 256)` is encoded by a
+**ResNet3D-50** backbone (Hara, Kataoka & Satoh, 2018; initialised from
+the Kinetics-700 release `r3d50_KM_200ep.pth` with conv1 weights
+summed across RGB → 1 grayscale channel).  Each of the four backbone
+stages is collapsed along the z (depth) axis with `torch.max`,
+producing a 2-D feature pyramid `{(256,64,64), (512,32,32),
+(1024,16,16), (2048,8,8)}`.  A small **2-D U-Net decoder** upsamples
+coarse-to-fine with concatenated skip connections; a 1×1 conv head
+produces a single sigmoid logit channel at quarter resolution
+`(B, 1, 64, 64)`.  Training uses `0.5·Dice + 0.5·SoftBCE` against the
+label down-interpolated to 64×64.
+
+## Quick start
+
+```python
+import torch
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained(
+    "YoussefMoNader/PHerc.1667-iteration-4",
+    trust_remote_code=True,
+).eval().cuda()
+
+# Input: float32, shape (B, 1, D=62, H=256, W=256).
+# Explicitly reproduce clipping followed by Albumentations'
+# Normalize(mean=0, std=1, max_pixel_value=255): clip(raw, 0, 200) / 255.
+# This is intensity scaling, not per-tile z-score normalization.
+# This synthetic tensor demonstrates the call shape only, not ink quality.
+raw = torch.randint(0, 256, (1, 1, 62, 256, 256), device="cuda")
+x = raw.float().clamp(0, 200) / 255.0
+
+with torch.no_grad():
+    out = model(x)
+
+print(out.logits.shape)                        # torch.Size([1, 1, 64, 64])
+prob = torch.sigmoid(out.logits)               # ink probability per pixel
+```
+
+## Full-segment inference (tiling)
+
+The model only sees 256×256 windows.  For a full scroll segment you
+need to slide the window across the (padded) layer stack and average
+overlapping predictions:
+
+```python
+import numpy as np, cv2, torch
+import torch.nn.functional as F
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained(
+    "YoussefMoNader/PHerc.1667-iteration-4", trust_remote_code=True,
+).eval().cuda()
+
+WINDOW, STRIDE = 256, 128       # 128 = 2x oversample; 64 for 8x oversample
+D = 62                           # number of z-layers
+
+# image: (H, W, D) uint8 stack of the 62 layers, padded to multiples of 256.
+# fmask: (H, W) uint8 fragment mask (0 = outside, 255 = inside).
+H, W, _ = image.shape
+mask_pred  = np.zeros((H, W), dtype=np.float32)
+mask_count = np.zeros((H, W), dtype=np.float32)
+
+with torch.no_grad():
+    for y in range(0, H - WINDOW + 1, STRIDE):
+        for x in range(0, W - WINDOW + 1, STRIDE):
+            if np.any(fmask[y:y+WINDOW, x:x+WINDOW] == 0):
+                continue
+            tile = image[y:y+WINDOW, x:x+WINDOW]            # (256,256,62)
+            # Apply the same input scaling as the quick start; keep image unchanged.
+            tile = np.clip(tile.astype(np.float32), 0, 200) / 255.0
+            t = torch.from_numpy(tile).permute(2, 0, 1)     # (62,256,256)
+            t = t.unsqueeze(0).unsqueeze(0).float().cuda()  # (1,1,62,256,256)
+            logits = model(t).logits                        # (1,1,64,64)
+            prob   = torch.sigmoid(logits)
+            prob   = F.interpolate(prob, scale_factor=4,
+                                   mode="bilinear").squeeze().cpu().numpy()
+            mask_pred[y:y+WINDOW, x:x+WINDOW]  += prob
+            mask_count[y:y+WINDOW, x:x+WINDOW] += 1.0
+
+pred = np.divide(mask_pred, mask_count,
+                 out=np.zeros_like(mask_pred),
+                 where=mask_count != 0)
+cv2.imwrite("prediction.png", np.clip(pred * 255, 0, 255).astype(np.uint8))
+```
+
+## Training summary
+
+| | |
+|---|---|
+| **Backbone** | ResNet3D-50 (3-D conv, BN, ReLU residual blocks) |
+| **Encoder init** | `r3d50_KM_200ep.pth` (Kinetics-700), conv1 summed across RGB |
+| **Decoder** | 2-D U-Net (3 up-blocks: bilinear 2× + concat skip + 3×3 conv + BN + ReLU) |
+| **Output** | 1 channel, sigmoid logit, quarter-resolution (64×64) |
+| **Loss** | 0.5 × Dice + 0.5 × SoftBCE (smooth = 0.25) |
+| **Optimizer** | AdamW, OneCycle lr 2e-5 → 3e-4, pct_start = 0.15 |
+| **Batch** | 2 (effective 8 via accumulate 4), 16-mixed, grad-clip 1.0 |
+| **Max steps** | 12,396 (= 3 epochs over the densest ablation label) |
+| **Training segment(s)** | `l_2` |
+| **Training label** | `l_2_inklabels4.png` |
+| **Training tiles** (256×256 sub-tiles at stride 64) | **24,773** |
+| **Final train loss (`_epoch`)** | **0.5256** |
+| **Final train loss (`_step`, single-batch noise)** | 0.5005 |
+| **Wandb** | [vesuvius-challenge/paper/l2_ink4_l5infer](https://wandb.ai/vesuvius-challenge/paper/runs/h73rxb10) |
+| **Random seed** | 130697 |
+| **Determinism** | `cudnn.deterministic = True`, `cudnn.benchmark = False` |
+| **Hardware** | 1 × NVIDIA H100 80 GB; ≈ 2 h end-to-end (load + train + inference) |
+
+## Files
+
+| file | size | description |
+|------|------|-------------|
+| `config.json` | 1 KB | architecture + provenance metadata; loaded by `AutoConfig` |
+| `configuration_inkdetection.py` | 2 KB | `InkDetectionConfig(PretrainedConfig)` |
+| `modeling_inkdetection.py` | 9 KB | self-contained `InkDetectionModel(PreTrainedModel)` |
+| `model.safetensors` | 319 MB | converted weights (338 tensors) |
+| `last.ckpt` | 963 MB | original PyTorch-Lightning checkpoint (incl. optimizer + LR-scheduler state) — load with `torch.load(...)["state_dict"]` |
+| `preview_l_2.png` | ~700 KB | low-res preview of the l_2 prediction (1/16 scale, 180° rotated) |
+| `preview_l_5.png` | ~2 MB | low-res preview of the l_5 (held-out) prediction |
+| `preview_label.png` | ~50 KB | the training label, same scale + rotation |
+
+The HuggingFace weights are **bit-perfect identical** to the original
+PyTorch-Lightning checkpoint (verified `max abs diff = 0.0e+00` on
+identical inputs).  Use `model.safetensors` for `AutoModel.from_pretrained`;
+use `last.ckpt` only if you want to resume training from the saved
+optimizer / scheduler state.
+
+## The model family
+
+| model | training segment(s) | label | tiles | effective epochs |
+|-------|---------------------|-------|-------|-------------------|
+| [`PHerc.1667-iteration-0`](https://huggingface.co/scrollprize/PHerc.1667-iteration-0) | 500p2a + 658 + 20250910185200 + 20250919125754* | (cross-segment baseline) | 20,075 | ~5 |
+| [`PHerc.1667-iteration-1`](https://huggingface.co/scrollprize/PHerc.1667-iteration-1) | `l_2` | `l_2_inklabels.png`  | 3,396 | ~30 |
+| [`PHerc.1667-iteration-2`](https://huggingface.co/scrollprize/PHerc.1667-iteration-2) | `l_2` | `l_2_inklabels2.png` | 8,970 | ~12 |
+| [`PHerc.1667-iteration-3`](https://huggingface.co/scrollprize/PHerc.1667-iteration-3) | `l_2` | `l_2_inklabels3.png` | 15,286 | ~7 |
+| [`PHerc.1667-iteration-4`](https://huggingface.co/scrollprize/PHerc.1667-iteration-4) | `l_2` | `l_2_inklabels4.png` | 24,773 | ~5 |
+| [`PHerc.1667-iteration-5`](https://huggingface.co/scrollprize/PHerc.1667-iteration-5) | `l_2` | `l_2_inklabels5.png` | 33,061 | 3 |
+
+All six share the architecture, hyperparameters, and a fixed step
+budget of 12,396 optimizer steps; the only thing that varies between
+rows is the supervising label (or, for ink0, the training segments).
+
+## Citation
+
+If you use this model in published work, please cite the Vesuvius
+Challenge and the underlying ResNet3D paper:
+
+```bibtex
+@inproceedings{hara2018can,
+  title  = {Can spatiotemporal 3D CNNs retrace the history of 2D CNNs and ImageNet?},
+  author = {Hara, Kensho and Kataoka, Hirokatsu and Satoh, Yutaka},
+  booktitle = {CVPR}, year = {2018},
+}
+```
+
+## Licence
+
+MIT.

modeling_inkdetection.py

@@ -1,239 +1,239 @@
-"""HuggingFace-style wrapper around the ResNet3D-50 + 2-D U-Net ink-detection model.
-
-This file is **self-contained** — vendored ResNet3D-50 (Hara et al., 2018)
-inline with the decoder so a downloader only needs `transformers` and
-`torch`.  Loadable via:
-
-    from transformers import AutoModel
-    model = AutoModel.from_pretrained(
-        "<user>/<repo>", trust_remote_code=True
-    )
-
-Input: float32 tensor of shape `(B, 1, D, H, W)` or `(B, D, H, W)`,
-       where D = 62, H = W = 256 (intensity already z-score normalised).
-Output: `ModelOutput(logits=<sigmoid logits, shape (B, 1, H/4, W/4)>)`.
-"""
-from dataclasses import dataclass
-from functools import partial
-from typing import List, Optional, Tuple
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from transformers import PreTrainedModel
-from transformers.modeling_outputs import ModelOutput
-
-from .configuration_inkdetection import InkDetectionConfig
-
-
-# =============================================================================
-# Vendored ResNet3D-50 (Hara, Kataoka & Satoh, 2018)
-# =============================================================================
-def _conv3x3x3(in_planes, out_planes, stride=1):
-    return nn.Conv3d(in_planes, out_planes, kernel_size=3,
-                     stride=stride, padding=1, bias=False)
-
-
-def _conv1x1x1(in_planes, out_planes, stride=1):
-    return nn.Conv3d(in_planes, out_planes, kernel_size=1,
-                     stride=stride, bias=False)
-
-
-class _Bottleneck(nn.Module):
-    expansion = 4
-
-    def __init__(self, in_planes, planes, stride=1, downsample=None):
-        super().__init__()
-        self.conv1 = _conv1x1x1(in_planes, planes)
-        self.bn1 = nn.BatchNorm3d(planes)
-        self.conv2 = _conv3x3x3(planes, planes, stride)
-        self.bn2 = nn.BatchNorm3d(planes)
-        self.conv3 = _conv1x1x1(planes, planes * self.expansion)
-        self.bn3 = nn.BatchNorm3d(planes * self.expansion)
-        self.relu = nn.ReLU(inplace=True)
-        self.downsample = downsample
-        self.stride = stride
-
-    def forward(self, x):
-        residual = x
-        out = self.relu(self.bn1(self.conv1(x)))
-        out = self.relu(self.bn2(self.conv2(out)))
-        out = self.bn3(self.conv3(out))
-        if self.downsample is not None:
-            residual = self.downsample(x)
-        out = self.relu(out + residual)
-        return out
-
-
-class _ResNet3D(nn.Module):
-    """ResNet3D-50 backbone returning the 4 intermediate feature maps."""
-
-    def __init__(self, n_input_channels=1, block_inplanes=(64, 128, 256, 512),
-                 layers=(3, 4, 6, 3), conv1_t_size=7, conv1_t_stride=1):
-        super().__init__()
-        self.in_planes = block_inplanes[0]
-        self.conv1 = nn.Conv3d(
-            n_input_channels, self.in_planes,
-            kernel_size=(conv1_t_size, 7, 7),
-            stride=(conv1_t_stride, 2, 2),
-            padding=(conv1_t_size // 2, 3, 3),
-            bias=False,
-        )
-        self.bn1 = nn.BatchNorm3d(self.in_planes)
-        self.relu = nn.ReLU(inplace=True)
-        self.maxpool = nn.MaxPool3d(
-            kernel_size=(1, 3, 3), stride=(1, 2, 2), padding=(0, 1, 1)
-        )
-        self.layer1 = self._make_layer(block_inplanes[0], layers[0], stride=1)
-        self.layer2 = self._make_layer(block_inplanes[1], layers[1], stride=2)
-        self.layer3 = self._make_layer(block_inplanes[2], layers[2], stride=2)
-        self.layer4 = self._make_layer(block_inplanes[3], layers[3], stride=2)
-
-    def _make_layer(self, planes, blocks, stride=1):
-        downsample = None
-        if stride != 1 or self.in_planes != planes * _Bottleneck.expansion:
-            downsample = nn.Sequential(
-                _conv1x1x1(self.in_planes,
-                           planes * _Bottleneck.expansion, stride),
-                nn.BatchNorm3d(planes * _Bottleneck.expansion),
-            )
-        layers = [_Bottleneck(self.in_planes, planes, stride, downsample)]
-        self.in_planes = planes * _Bottleneck.expansion
-        for _ in range(1, blocks):
-            layers.append(_Bottleneck(self.in_planes, planes))
-        return nn.Sequential(*layers)
-
-    def forward(self, x) -> List[torch.Tensor]:
-        x = self.relu(self.bn1(self.conv1(x)))
-        x = self.maxpool(x)
-        x1 = self.layer1(x)
-        x2 = self.layer2(x1)
-        x3 = self.layer3(x2)
-        x4 = self.layer4(x3)
-        return [x1, x2, x3, x4]
-
-
-# =============================================================================
-# 2-D U-Net decoder
-# =============================================================================
-class _Decoder(nn.Module):
-    def __init__(self, encoder_dims, upscale):
-        super().__init__()
-        self.convs = nn.ModuleList([
-            nn.Sequential(
-                nn.Conv2d(encoder_dims[i] + encoder_dims[i - 1],
-                          encoder_dims[i - 1], 3, 1, 1, bias=False),
-                nn.BatchNorm2d(encoder_dims[i - 1]),
-                nn.ReLU(inplace=True),
-            )
-            for i in range(1, len(encoder_dims))
-        ])
-        self.logit = nn.Conv2d(encoder_dims[0], 1, 1, 1, 0)
-        self.up = nn.Upsample(scale_factor=upscale, mode='bilinear')
-
-    def forward(self, feature_maps):
-        for i in range(len(feature_maps) - 1, 0, -1):
-            f_up = F.interpolate(feature_maps[i], scale_factor=2,
-                                 mode='bilinear')
-            f = torch.cat([feature_maps[i - 1], f_up], dim=1)
-            feature_maps[i - 1] = self.convs[i - 1](f)
-        return self.up(self.logit(feature_maps[0]))
-
-
-# =============================================================================
-# HuggingFace ModelOutput + PreTrainedModel
-# =============================================================================
-@dataclass
-class InkDetectionOutput(ModelOutput):
-    """Output of `InkDetectionModel.forward`.
-
-    - `logits`: pre-sigmoid prediction at quarter resolution.  Shape
-      `(B, 1, H / 4, W / 4)`. Apply `torch.sigmoid` for probabilities.
-    """
-    logits: torch.FloatTensor = None
-    loss: Optional[torch.FloatTensor] = None
-
-
-class InkDetectionModel(PreTrainedModel):
-    """Vesuvius Challenge ink-detection model.
-
-    Pipeline:
-        1.  3-D volume `(B, 1, D, H, W)` enters a ResNet3D-50 backbone.
-        2.  Each of the 4 stages is collapsed along the z (depth) axis
-            with `torch.max` -> 2-D feature pyramid.
-        3.  A small 2-D U-Net decoder upsamples coarse-to-fine with
-            concatenated skip connections.
-        4.  A 1x1 conv head produces 1 logit channel.
-    """
-    config_class = InkDetectionConfig
-    base_model_prefix = "inkdetection"
-
-    def __init__(self, config: InkDetectionConfig):
-        super().__init__(config)
-        layers_map = {50: (3, 4, 6, 3),
-                      101: (3, 4, 23, 3),
-                      152: (3, 8, 36, 3)}
-        if config.backbone_depth not in layers_map:
-            raise ValueError(
-                f"Unsupported backbone_depth={config.backbone_depth}; "
-                "expected one of 50, 101, 152."
-            )
-        self.backbone = _ResNet3D(
-            n_input_channels=config.in_channels,
-            block_inplanes=(64, 128, 256, 512),
-            layers=layers_map[config.backbone_depth],
-        )
-        self.decoder = _Decoder(
-            encoder_dims=list(config.backbone_channels),
-            upscale=config.decoder_upscale,
-        )
-        # No init from scratch — weights are loaded from the published
-        # checkpoint via `from_pretrained`.
-
-    def forward(
-        self,
-        pixel_values: torch.FloatTensor,
-        labels: Optional[torch.FloatTensor] = None,
-        return_dict: Optional[bool] = True,
-        **kwargs,
-    ) -> InkDetectionOutput:
-        # Accept (B, D, H, W) or (B, 1, D, H, W)
-        if pixel_values.ndim == 4:
-            pixel_values = pixel_values.unsqueeze(1)
-        if pixel_values.ndim != 5:
-            raise ValueError(
-                f"pixel_values must be 4-D (B, D, H, W) or 5-D (B, 1, D, H, W); "
-                f"got shape {tuple(pixel_values.shape)}"
-            )
-        feats = self.backbone(pixel_values)
-        pooled = [torch.max(f, dim=2)[0] for f in feats]
-        logits = self.decoder(pooled)
-
-        loss = None
-        if labels is not None:
-            # Dice + SoftBCE in [0, 1] target space, label assumed to
-            # already be down-interpolated to logits.shape[-2:].
-            sig = torch.sigmoid(logits)
-            inter = (sig * labels).sum(dim=(-2, -1))
-            denom = sig.sum(dim=(-2, -1)) + labels.sum(dim=(-2, -1))
-            dice = (1.0 - (2.0 * inter + 1.0) / (denom + 1.0)).mean()
-            bce = F.binary_cross_entropy_with_logits(logits, labels)
-            loss = 0.5 * dice + 0.5 * bce
-
-        if not return_dict:
-            return (loss, logits) if loss is not None else (logits,)
-        return InkDetectionOutput(logits=logits, loss=loss)
-
-    @torch.no_grad()
-    def predict_probability(self, pixel_values: torch.FloatTensor) -> torch.FloatTensor:
-        """Convenience: sigmoid probabilities at quarter resolution."""
-        return torch.sigmoid(self(pixel_values).logits)
-
-
-# Register the config so HF's mapping infrastructure recognises model_type.
-# Note: AutoModel.from_pretrained(..., trust_remote_code=True) reads `auto_map`
-# from config.json — no explicit register call is required at import time, but
-# it does not hurt to keep this association in place.
-InkDetectionConfig.register_for_auto_class("AutoConfig")
-InkDetectionModel.register_for_auto_class("AutoModel")
+"""HuggingFace-style wrapper around the ResNet3D-50 + 2-D U-Net ink-detection model.
+
+This file is **self-contained** — vendored ResNet3D-50 (Hara et al., 2018)
+inline with the decoder so a downloader only needs `transformers` and
+`torch`.  Loadable via:
+
+    from transformers import AutoModel
+    model = AutoModel.from_pretrained(
+        "<user>/<repo>", trust_remote_code=True
+    )
+
+Input: float32 tensor of shape `(B, 1, D, H, W)` or `(B, D, H, W)`,
+       where D = 62, H = W = 256 (raw uint8 clipped to [0, 200] and divided by 255; not z-score normalised).
+Output: `ModelOutput(logits=<sigmoid logits, shape (B, 1, H/4, W/4)>)`.
+"""
+from dataclasses import dataclass
+from functools import partial
+from typing import List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+
+from .configuration_inkdetection import InkDetectionConfig
+
+
+# =============================================================================
+# Vendored ResNet3D-50 (Hara, Kataoka & Satoh, 2018)
+# =============================================================================
+def _conv3x3x3(in_planes, out_planes, stride=1):
+    return nn.Conv3d(in_planes, out_planes, kernel_size=3,
+                     stride=stride, padding=1, bias=False)
+
+
+def _conv1x1x1(in_planes, out_planes, stride=1):
+    return nn.Conv3d(in_planes, out_planes, kernel_size=1,
+                     stride=stride, bias=False)
+
+
+class _Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, in_planes, planes, stride=1, downsample=None):
+        super().__init__()
+        self.conv1 = _conv1x1x1(in_planes, planes)
+        self.bn1 = nn.BatchNorm3d(planes)
+        self.conv2 = _conv3x3x3(planes, planes, stride)
+        self.bn2 = nn.BatchNorm3d(planes)
+        self.conv3 = _conv1x1x1(planes, planes * self.expansion)
+        self.bn3 = nn.BatchNorm3d(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+        out = self.relu(self.bn1(self.conv1(x)))
+        out = self.relu(self.bn2(self.conv2(out)))
+        out = self.bn3(self.conv3(out))
+        if self.downsample is not None:
+            residual = self.downsample(x)
+        out = self.relu(out + residual)
+        return out
+
+
+class _ResNet3D(nn.Module):
+    """ResNet3D-50 backbone returning the 4 intermediate feature maps."""
+
+    def __init__(self, n_input_channels=1, block_inplanes=(64, 128, 256, 512),
+                 layers=(3, 4, 6, 3), conv1_t_size=7, conv1_t_stride=1):
+        super().__init__()
+        self.in_planes = block_inplanes[0]
+        self.conv1 = nn.Conv3d(
+            n_input_channels, self.in_planes,
+            kernel_size=(conv1_t_size, 7, 7),
+            stride=(conv1_t_stride, 2, 2),
+            padding=(conv1_t_size // 2, 3, 3),
+            bias=False,
+        )
+        self.bn1 = nn.BatchNorm3d(self.in_planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool3d(
+            kernel_size=(1, 3, 3), stride=(1, 2, 2), padding=(0, 1, 1)
+        )
+        self.layer1 = self._make_layer(block_inplanes[0], layers[0], stride=1)
+        self.layer2 = self._make_layer(block_inplanes[1], layers[1], stride=2)
+        self.layer3 = self._make_layer(block_inplanes[2], layers[2], stride=2)
+        self.layer4 = self._make_layer(block_inplanes[3], layers[3], stride=2)
+
+    def _make_layer(self, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.in_planes != planes * _Bottleneck.expansion:
+            downsample = nn.Sequential(
+                _conv1x1x1(self.in_planes,
+                           planes * _Bottleneck.expansion, stride),
+                nn.BatchNorm3d(planes * _Bottleneck.expansion),
+            )
+        layers = [_Bottleneck(self.in_planes, planes, stride, downsample)]
+        self.in_planes = planes * _Bottleneck.expansion
+        for _ in range(1, blocks):
+            layers.append(_Bottleneck(self.in_planes, planes))
+        return nn.Sequential(*layers)
+
+    def forward(self, x) -> List[torch.Tensor]:
+        x = self.relu(self.bn1(self.conv1(x)))
+        x = self.maxpool(x)
+        x1 = self.layer1(x)
+        x2 = self.layer2(x1)
+        x3 = self.layer3(x2)
+        x4 = self.layer4(x3)
+        return [x1, x2, x3, x4]
+
+
+# =============================================================================
+# 2-D U-Net decoder
+# =============================================================================
+class _Decoder(nn.Module):
+    def __init__(self, encoder_dims, upscale):
+        super().__init__()
+        self.convs = nn.ModuleList([
+            nn.Sequential(
+                nn.Conv2d(encoder_dims[i] + encoder_dims[i - 1],
+                          encoder_dims[i - 1], 3, 1, 1, bias=False),
+                nn.BatchNorm2d(encoder_dims[i - 1]),
+                nn.ReLU(inplace=True),
+            )
+            for i in range(1, len(encoder_dims))
+        ])
+        self.logit = nn.Conv2d(encoder_dims[0], 1, 1, 1, 0)
+        self.up = nn.Upsample(scale_factor=upscale, mode='bilinear')
+
+    def forward(self, feature_maps):
+        for i in range(len(feature_maps) - 1, 0, -1):
+            f_up = F.interpolate(feature_maps[i], scale_factor=2,
+                                 mode='bilinear')
+            f = torch.cat([feature_maps[i - 1], f_up], dim=1)
+            feature_maps[i - 1] = self.convs[i - 1](f)
+        return self.up(self.logit(feature_maps[0]))
+
+
+# =============================================================================
+# HuggingFace ModelOutput + PreTrainedModel
+# =============================================================================
+@dataclass
+class InkDetectionOutput(ModelOutput):
+    """Output of `InkDetectionModel.forward`.
+
+    - `logits`: pre-sigmoid prediction at quarter resolution.  Shape
+      `(B, 1, H / 4, W / 4)`. Apply `torch.sigmoid` for probabilities.
+    """
+    logits: torch.FloatTensor = None
+    loss: Optional[torch.FloatTensor] = None
+
+
+class InkDetectionModel(PreTrainedModel):
+    """Vesuvius Challenge ink-detection model.
+
+    Pipeline:
+        1.  3-D volume `(B, 1, D, H, W)` enters a ResNet3D-50 backbone.
+        2.  Each of the 4 stages is collapsed along the z (depth) axis
+            with `torch.max` -> 2-D feature pyramid.
+        3.  A small 2-D U-Net decoder upsamples coarse-to-fine with
+            concatenated skip connections.
+        4.  A 1x1 conv head produces 1 logit channel.
+    """
+    config_class = InkDetectionConfig
+    base_model_prefix = "inkdetection"
+
+    def __init__(self, config: InkDetectionConfig):
+        super().__init__(config)
+        layers_map = {50: (3, 4, 6, 3),
+                      101: (3, 4, 23, 3),
+                      152: (3, 8, 36, 3)}
+        if config.backbone_depth not in layers_map:
+            raise ValueError(
+                f"Unsupported backbone_depth={config.backbone_depth}; "
+                "expected one of 50, 101, 152."
+            )
+        self.backbone = _ResNet3D(
+            n_input_channels=config.in_channels,
+            block_inplanes=(64, 128, 256, 512),
+            layers=layers_map[config.backbone_depth],
+        )
+        self.decoder = _Decoder(
+            encoder_dims=list(config.backbone_channels),
+            upscale=config.decoder_upscale,
+        )
+        # No init from scratch — weights are loaded from the published
+        # checkpoint via `from_pretrained`.
+
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor,
+        labels: Optional[torch.FloatTensor] = None,
+        return_dict: Optional[bool] = True,
+        **kwargs,
+    ) -> InkDetectionOutput:
+        # Accept (B, D, H, W) or (B, 1, D, H, W)
+        if pixel_values.ndim == 4:
+            pixel_values = pixel_values.unsqueeze(1)
+        if pixel_values.ndim != 5:
+            raise ValueError(
+                f"pixel_values must be 4-D (B, D, H, W) or 5-D (B, 1, D, H, W); "
+                f"got shape {tuple(pixel_values.shape)}"
+            )
+        feats = self.backbone(pixel_values)
+        pooled = [torch.max(f, dim=2)[0] for f in feats]
+        logits = self.decoder(pooled)
+
+        loss = None
+        if labels is not None:
+            # Dice + SoftBCE in [0, 1] target space, label assumed to
+            # already be down-interpolated to logits.shape[-2:].
+            sig = torch.sigmoid(logits)
+            inter = (sig * labels).sum(dim=(-2, -1))
+            denom = sig.sum(dim=(-2, -1)) + labels.sum(dim=(-2, -1))
+            dice = (1.0 - (2.0 * inter + 1.0) / (denom + 1.0)).mean()
+            bce = F.binary_cross_entropy_with_logits(logits, labels)
+            loss = 0.5 * dice + 0.5 * bce
+
+        if not return_dict:
+            return (loss, logits) if loss is not None else (logits,)
+        return InkDetectionOutput(logits=logits, loss=loss)
+
+    @torch.no_grad()
+    def predict_probability(self, pixel_values: torch.FloatTensor) -> torch.FloatTensor:
+        """Convenience: sigmoid probabilities at quarter resolution."""
+        return torch.sigmoid(self(pixel_values).logits)
+
+
+# Register the config so HF's mapping infrastructure recognises model_type.
+# Note: AutoModel.from_pretrained(..., trust_remote_code=True) reads `auto_map`
+# from config.json — no explicit register call is required at import time, but
+# it does not hurt to keep this association in place.
+InkDetectionConfig.register_for_auto_class("AutoConfig")
+InkDetectionModel.register_for_auto_class("AutoModel")