RockSinGAN

This repository contains the RockSinGAN code with pytorch framework. This work is related to the article "RockSinGAN: Single-Image Generative AI for Multiscale 3D Rock generation", which was exclusively submitted to the Journal of Computers & Geosciences.

RockSinGAN is a 3D generative framework designed to generate digital rock models based in a training image. Although our application focused on the digital rock physics area, the RockSinGAN framework is a versatile tool that can be applied to other geoscientific problems as a robust geostatistical simulation technique.

Highlights

• Hierarchical 3D Generation: A robust geostatistical simulation technique trained on a single image to learn complex rock textures.
• Memory-Efficient Architecture: Combines Patched Inference and Early Cropping to generate ultra-high-resolution models (up to $3344 \times 1360 \times 8304$) on standard hardware.
• Data integration: Application to integrate coreCT and microCT. It integrates micro-scale pore features into large-scale rock models, overcoming the traditional trade-off between sample size and image resolution.
• Hard Data Conditioning: A critical feature for geological applications that enables the seamless incorporation of direct measurements into generated models[cite: 50, 317, 878]. This scheme preserves well-resolved macropores and matrix from low-resolution coreCT while populating unresolved regions with high-resolution features, ensuring spatial consistency across scales.
• High Efficiency: The Early Cropping approach completes the generation process in just 64.8% of the time required by standard patched inference.

Structure

Main folders and files in rocksingan

• ./config/: Program configuration variables.
• ./model/: GAN model files.
• ./finetune.py: Script to fine-tunine a trained model.
• ./resume.py: Script to resume an incomplete trained model.
• ./train.py: Main file to run the training.
• ./image_generation.py: Script to load a trained model and generate samples.
• ./noise_generation.py: Script to create and save on disk the noise required to generate a sample.

How to setup?

We recommend installing the project in a virtual environment with Python>=3.9.

First, make sure you have the Pytorch with the correct cuda installed in your environment. Example:

pip install torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 --index-url https://download.pytorch.org/whl/cu126

Install this library via pip in your Python env:

pip install <path-to-this-folder>

or 

python3 -m pip install <path-to-this-folder>

Usage

Training

To train the GAN use the tools inside the train.py file:

Running via command line

# Check the configuration parameters
python3 -m rocksingan.train -h
python3 -m rocksingan.train --input_path <path-to-input-image> ...

Running inside a project

from rocksingan.config.config import Opt
from rocksingan.train import train

opt = Opt(
    ti_resolution_mm=0.2,
    input_path="<path-to-input-image>",
    min_size=32,
    max_size=128,
    crop_size=64,
    img_color_range=[0, 255],
    out_path="<output-path-name>",
    num_iter=2000,
)

train(opt)

Resume Training

If you need to resume an incomplete training, use the tools inside the resume.py file:

Running via command line

python3 -m rocksingan.resume --checkpoint_path <path-of-the-model>

Running inside a project

from rocksingan.resume import resume_train

resume_training("<path-of-the-model>")

Finetune a Trained Model

If you want to finetune a trained model, use the tools inside the finetune.py file:

Running via command line

python3 -m rocksingan.finetune --checkpoint_path <path-of-the-model> --start_scale 0 --num_iter 1000

Running inside a project

from rocksingan.finetuning import finetune

finetune(checkpoint_path="<path-of-the-model>", start_scale=0, num_iter=1000)

Generation

You can use the tools inside the image_generation.py file to generate samples:

Running via command line

python3 -m rocksingan.image_generation --model_path <path-of-the-model> ...

Running inside a project

from rocksingan.image_generation import generate_image

generate_image(
    model_path="2025_09_22_11_36_28_teste", 
    technique="patched",
    out_dir="./",
    base_volume=50,
    disk_scale=0,
    split_scale=0,
    gpu_device=0,
    noise_type="rec",
    noise_seed=None,
    partitions=1,
    corect_path=None,
    corect_resolution=None,
    injection_start_scale=None,
    auto_params=False,
    hard_data=None,
)

Parameters description

Global params:

• model_path: path to the trained model weights;
• technique: "patched" or "early_crop";
• out_dir: the folder where the sample will be saved;
• partitions: the number of .nc files generated. For example, if the sample size is 1000x1000x1000 and "partitions" is set to 2, the sample will be divided into two sections along the first axis, and two .nc files will be generated. Each .nc file corresponds to a 500x1000x1000 section of the sample;
• gpu_device: GPU device index;

Generation params:

• base_volume: the size of the patch processed by the GPU;
• disk_scale: the scale from which intermediate data will be saved in disk;
• split_scale: (required by early crop technique) the scale from which the upscaling of an output scale will be made in chunks;
• auto_params: if True, "base_volume", "disk_scale", and "split_scale" will be estimated based on the currently available computational resources;

Noise params:

• noise_type: "rec", "dynamic", or a path to a noise saved on disk;

• noise_seed: (optional) a seed used to generate the dynamic noise;

  Note: Dynamic noise depends on "base_volume" and "noise_seed". To reproduce the same noise across different generations, these two parameters must therefore be fixed. If you want to fix the noise but change the “base_volume” parameter, use the noise_generation.py script to create and save on disk the noise before starting image generation.

coreCT integration params:

• corect_path: path to the coreCT image;
• corect_resolution: coreCT image resolution (corect_resolution or injection_start_scale must be specified)
• injection_start_scale: the scale from which the coreCT injection will be made (corect_resolution or injection_start_scale must be specified)
• hard_data: target class/classes that will be used as hard data.

Noise generation

To create and save on disk noise before starting image generation:

python3 -m rocksingan.noise_generation --model_path <path-of-the-model> ...

Parameters:

• --model_path: path to the trained model weights;
• --out_dir: the folder where the noise will be saved;
• --folder_name: noise folder name;
• --seed: (optional) seed used to generate the noise;
• --corect_path: path to the coreCT image;
• --corect_resolution: coreCT image resolution (corect_resolution or injection_start_scale must be specified)
• --injection_start_scale: the scale from which the coreCT injection will be made (corect_resolution or injection_start_scale must be specified)

Reference

• SinGAN: Learning a Generative Model from a Single Natural Image
• BlendGAN: Learning and Blending the Internal Distributions of Single Images by Spatial Image-Identity Conditioning
• tamarott/SinGAN
• tantelu/CSinGAN