mwalmsley / euclid-rr2-mae

Masked Autoencoder for Euclid Images

Overview

This masked autoencoder (MAE) is trained to reconstruct Euclid galaxy images where 90% of the image is masked. The trained model shows superhuman performance at reconstruction. Examples are shown in Galaxy Morphology and Interpretability through Sparsity, Wu & Walmsley, NeurIPS ML4Science workshop 2025 (on arxiv very shortly, bibtex to follow).

If I get time, I'll make an interactive demo.

This version is trained on RR2 (3M images). A DR1 version (13.6M images) will follow.

Instructions

Download the model

import mae_timm_simplified  # download this script from the "files and versions" tab

import omegaconf
from huggingface_hub import hf_hub_download

cfg_path = hf_hub_download(repo_id="mwalmsley/euclid-rr2-mae", filename="config.yaml")
cfg = omegaconf.OmegaConf.load(cfg_path)
mae = mae_timm_simplified.MAE.from_pretrained("mwalmsley/euclid-rr2-mae", cfg=cfg)

Download Data

Get a dataset of Euclid images, prepared as Galaxy-Zoo-style jpgs:

from datasets import load_dataset

dataset_dict = load_dataset(
    'mwalmsley/euclid_q1',   # _rr2, _dr1 versions are available to EC members
    name='tiny-v1-gz_arcsinh_vis_y'  # tiny subset for testing
)

Use my utility package galaxy-datasets to load this as a Lightning DataModule, including an appropriate torchvision transform...

from galaxy_datasets.pytorch.galaxy_datamodule import HuggingFaceDataModule
from galaxy_datasets.transforms import default_view_config, get_galaxy_transform


# define augmentations to use
view_config = default_view_config()
view_config.output_size = 224
view_config.erase_iterations = 0  # for simplicity
ssl_image_transform = get_galaxy_transform(cfg=view_config)
# this is just a torchvision Compose transform
# returns 3x224x224 float tensor normalised 0-1.

datamodule = HuggingFaceDataModule(
    dataset_dict=dataset_dict,
    train_transform=ssl_image_transform,
    test_transform=ssl_image_transform,
    batch_size=batch_size,
    num_workers=num_workers,
    prefetch_factor=prefetch_factor
)
datamodule.setup()
# this is just a lightning datamodule
# should yield batches with an 'image' key, see below

# get a batch
test_loader = datamodule.test_dataloader()
for batch in test_loader:
    batch['image'] = batch['image'].to('cuda')
    break

...or you can do this yourself. You should make batches that include an 'image' key which contains

• BxCx224x224 float tensors normalised from 0 to 1
• where those tensors are created by transforming (e.g. with torchvision) a GZ-style jpg (download from HuggingFace above)

It might work for other human-friendly jpgs, but that's outside of the training distribution, so no promises.

For example:

(TODO)

Mask some pixels

Generate a random mask (as during training)

    
    batch_size = batch['image'].shape[0]
    _, idx_mask = random_token_mask(
        size=(batch_size, mae.sequence_length),  # (batch_size, seq_len)
        mask_ratio=0.9,
        device='cuda'
    )

Generate a custom mask

We divide the image into 784 patches, in a grid of 28 by 28 patches. Each patch is 8x8 pixels (covering our 224x224 image).

The patch in the top corner is index 1 (not 0!), and higher indices go left-to-right and then down a row (like reading a page).

For example, to mask only the first 28 patches (the top of the image)

    row_mask = torch.tensor(range(28)) + 1 # 1 to 29
    # copy for all images in the batch
    idx_mask = row_mask.unsqueeze(0).repeat(batch_size, 1)  # (batch_size, num_masked)
    idx_mask = idx_mask.to('cuda')

To mask the middle strip:

    row_mask = torch.tensor(range(28)) + 1 + 13*28

And so on, however you like. Just remember to add 1 for the class token!

Make predictions for the masked patches

    mae = mae.to('cuda')
    with torch.no_grad():
        result = mae.predict(batch, idx_mask=idx_mask)

    # result has keys including images, masked, reconstructed
    # each key is a list of standard PIL images
    images = result['images']
    masked = result['masked']
    reconstructed = result['reconstructed']

    # Visualize the results
    fig, axes = plt.subplots(nrows=3, ncols=8, figsize=(24, 9))
    for i in range(8):
        axes[0, i].imshow(images[i])
        axes[0, i].set_title("Original")
        axes[1, i].imshow(masked[i])
        axes[1, i].set_title("Masked")
        axes[2, i].imshow(reconstructed[i])
        axes[2, i].set_title("Reconstructed")
    plt.tight_layout()
    plt.show()