Initial commit: weights + custom modeling

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+preview_l_2.png filter=lfs diff=lfs merge=lfs -text
+preview_l_5.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,193 @@
+---
+library_name: transformers
+license: mit
+tags:
+- vesuvius-challenge
+- ink-detection
+- herculaneum
+- resnet3d
+- u-net
+- 3d-segmentation
+- volumetric-imaging
+pipeline_tag: image-segmentation
+---
+
+# vesuvius-inkdetection-ink1
+
+> **Trained on segment l_2 with l_2_inklabels.png (3,396 tiles).**
+
+Ablation 1/5 — sparsest training label (3,396 tiles). Smallest annotation set; trained for ~30 effective epochs over its data to match the step budget.
+
+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/vesuvius-inkdetection-ink1",
+    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/vesuvius-inkdetection-ink1", 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_inklabels.png` |
+| **Training tiles** (256×256 sub-tiles at stride 64) | **3,396** |
+| **Final train loss (`_epoch`)** | **0.4219** |
+| **Final train loss (`_step`, single-batch noise)** | 0.4381 |
+| **Wandb** | [vesuvius-challenge/Nature/l2_ink1_l5infer](https://wandb.ai/Nature/l2_ink1_l5infer) |
+| **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 |
+|-------|---------------------|-------|-------|-------------------|
+| [`vesuvius-inkdetection-ink0`](https://huggingface.co/YoussefMoNader/vesuvius-inkdetection-ink0) | 500p2a + 658 + 20250910185200 + 20250919125754* | (cross-segment baseline) | 20,075 | ~5 |
+| [`vesuvius-inkdetection-ink1`](https://huggingface.co/YoussefMoNader/vesuvius-inkdetection-ink1) | `l_2` | `l_2_inklabels.png`  | 3,396 | ~30 |
+| [`vesuvius-inkdetection-ink2`](https://huggingface.co/YoussefMoNader/vesuvius-inkdetection-ink2) | `l_2` | `l_2_inklabels2.png` | 8,970 | ~12 |
+| [`vesuvius-inkdetection-ink3`](https://huggingface.co/YoussefMoNader/vesuvius-inkdetection-ink3) | `l_2` | `l_2_inklabels3.png` | 15,286 | ~7 |
+| [`vesuvius-inkdetection-ink4`](https://huggingface.co/YoussefMoNader/vesuvius-inkdetection-ink4) | `l_2` | `l_2_inklabels4.png` | 24,773 | ~5 |
+| [`vesuvius-inkdetection-ink5`](https://huggingface.co/YoussefMoNader/vesuvius-inkdetection-ink5) | `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.

config.json

@@ -0,0 +1,29 @@
+{
+  "model_type": "inkdetection_resnet3d",
+  "architectures": [
+    "InkDetectionModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_inkdetection.InkDetectionConfig",
+    "AutoModel": "modeling_inkdetection.InkDetectionModel"
+  },
+  "in_channels": 1,
+  "input_depth": 62,
+  "input_size": 256,
+  "backbone_depth": 50,
+  "backbone_channels": [
+    256,
+    512,
+    1024,
+    2048
+  ],
+  "num_classes": 1,
+  "decoder_upscale": 1,
+  "train_segment": "l_2",
+  "train_inklabels": "l_2_inklabels.png",
+  "train_steps": 12396,
+  "train_tiles": 3396,
+  "train_loss_final": 0.4219,
+  "torch_dtype": "float32",
+  "transformers_version": "4.57.6"
+}

configuration_inkdetection.py

@@ -0,0 +1,48 @@
+"""Configuration class for the Vesuvius ink-detection model.
+
+See modeling_inkdetection.py for the model itself. Designed to be
+loaded with `AutoConfig.from_pretrained(..., trust_remote_code=True)`.
+"""
+from transformers import PretrainedConfig
+
+
+class InkDetectionConfig(PretrainedConfig):
+    """Configuration for an InkDetectionModel.
+
+    Architecture: ResNet3D-50 backbone (Hara et al., 2018) initialised
+    from Kinetics-700; per-stage max-pool over the z (depth) axis to
+    collapse 5-D features to 4-D; small 2-D U-Net decoder; 1x1 conv
+    output head producing one sigmoid logit channel.
+    """
+    model_type = "inkdetection_resnet3d"
+
+    def __init__(
+        self,
+        in_channels: int = 1,
+        input_depth: int = 62,
+        input_size: int = 256,
+        backbone_depth: int = 50,
+        backbone_channels=(256, 512, 1024, 2048),
+        num_classes: int = 1,
+        decoder_upscale: int = 1,
+        # Optional metadata (informational only)
+        train_segment: str = "l_2",
+        train_inklabels: str = "l_2_inklabels.png",
+        train_steps: int = 12396,
+        train_tiles: int = 0,
+        train_loss_final: float = 0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.input_depth = input_depth
+        self.input_size = input_size
+        self.backbone_depth = backbone_depth
+        self.backbone_channels = list(backbone_channels)
+        self.num_classes = num_classes
+        self.decoder_upscale = decoder_upscale
+        self.train_segment = train_segment
+        self.train_inklabels = train_inklabels
+        self.train_steps = train_steps
+        self.train_tiles = train_tiles
+        self.train_loss_final = train_loss_final

last.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:04274e174463a7feecac912abab2ef62844cecbf93e5de25d3f1e2b971938d0b
+size 1008820953

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:75e4009673b2c419c23cb388474c87357c226058598f1e2d7cba4e71d22e0fcb
+size 333531996

modeling_inkdetection.py

@@ -0,0 +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")