Input and output formats

• 2D inputs: .png, .jpg, or .jpeg, converted to grayscale.

• 3D inputs: .mat files containing a three-dimensional array named temp.

• Image intensities are expected on a 0–255 scale and normalized to [-1, 1].

• Use .tif for image or volume output, or .pt to save the prediction as a PyTorch tensor. TIFF output is converted to unsigned 8-bit intensities; tensor output retains the model’s values.

Prediction options

The current prediction command resizes the input to the requested spatial size before running the model. For example, --size 100 resizes a 3D volume to 100 × 100 × 100 voxels. Choose a size compatible with your model and intended output resolution. The training option --scale is not a prediction option. Diffusion sampling options apply when a diffusion checkpoint is loaded.

Dataset layout

Pass the directory containing the rock category folders to --deeprock. The loaders search for sandstone, carbonate, coal, and sand; provide training and validation pairs for the categories you use.

For 3D data at scale 4, the layout for one category is:

deeprock/
    sandstone3D/
        sandstone3D_train_HR/
            sample.mat
        sandstone3D_train_LR_default_X4/
            samplex4.mat
        sandstone3D_valid_HR/
            sample.mat
        sandstone3D_valid_LR_default_X4/
            samplex4.mat

Each MATLAB file must contain an array named temp. The 3D loader upsamples the low-resolution volume by --scale using cubic interpolation before passing it to the model; the resulting shape must match the high-resolution volume.

For 2D data at scale 4, the layout is:

deeprock/
    sandstone2D/
        sandstone2D_train_HR/
            sample.png
        sandstone2D_train_BI_unknown_X4/
            sample.png
        sandstone2D_valid_HR/
            sample.png
        sandstone2D_valid_BI_unknown_X4/
            sample.png

The 2D loader uses the same filename for each pair and reads the images at their stored sizes. Prepare the low-resolution inputs at the spatial size expected by your model. For other scales, replace X4 in directory names and x4 in 3D filenames with the selected scale.

Training options

Use supercat-tools train --help for the full option list, including model presets, U-Net selection, transformer dimensions, and checkpoint options.

Image pretraining

Train a 2D model from directories of PNG or JPEG images:

supercat-pretrain-image train \
    --training /path/to/images/train \
    --validation /path/to/images/valid \
    --dim 2 \
    --scale 4 \
    --min-size 224 \
    --max-size 224 \
    --epochs 40 \
    --results-dir results/pretrain-image

The image dataset uses bicubic interpolation to create the degraded input. Setting --min-size and --max-size to the same value produces samples with a consistent spatial size; use an even size compatible with your model.

Video pretraining

Train a 3D model using clips from video files:

supercat-pretrain-movie train \
    --training /path/to/videos/train \
    --validation /path/to/videos/valid \
    --dim 3 \
    --scale 4 \
    --min-size 100 \
    --max-size 100 \
    --max-training-items 1000 \
    --max-validation-items 100 \
    --epochs 40 \
    --results-dir results/pretrain-movie

The video dataset discovers .mp4, .avi, .mov, and .mkv files. It treats the frame axis as the third spatial dimension and uses trilinear interpolation to create the degraded input. A working video decoding backend is required, such as imageio[ffmpeg] or an ffmpeg executable on PATH.

Pretraining options

For images, size limits apply to height and width. For videos, they also apply to the frame count. Validation uses centered crops without augmentation. The item limits count video files, not individual frames.

Both commands also accept the model and training options described above, including --no-use-diffusion for regression training. Inspect the complete options with:

supercat-pretrain-image train --help
supercat-pretrain-movie train --help