XAMI-model: XMM-Newton optical Artefact Mapping for astronomical Instance segmentation

This repository contains the weights of the XAMI model. The model is trained on images from the XAMI dataset (available on Github and HuggingFace). The images are astronomical observations from the Optical Monitor (XMM-OM) onboard the XMM-Newton X-ray mission of the European Space Agency (ESA).

Information about the XMM-OM can be found here:

• The ESA website
• The article The XMM-Newton optical/UV monitor telescope.

Cloning the repository

git clone https://github.com/ESA-Datalabs/XAMI-model.git
cd XAMI-model

# creating the environment
conda env create -f environment.yaml
conda activate xami_model_env

# Install the package in editable mode
pip install -e .

Downloading the dataset and model checkpoints from HuggingFace

The dataset is splited into train and validation categories and contains annotated artefacts in COCO format for Instance Segmentation. We use multilabel Stratified K-fold (k=4) to balance class distributions across splits. We choose to work with a single dataset splits version (out of 4) but also provide means to work with all 4 versions.

The Dataset-Structure.md offers more details about the dataset structure. We provide the following dataset formats: COCO format for Instance Segmentation (commonly used by Detectron2 models) and YOLOv8-Seg format used by ultralytics.

Check the dataset_and_model.ipynb for downloading the dataset and model weights.

Model Inference

After cloning the repository and setting up the environment, use the following code for model loading and inference:

from xami_model.inference.xami_inference import InferXami

det_type = 'rtdetr' # 'rtdetr'  'yolov8'

detr_checkpoint = f'./xami_model/train/weights/{det_type}_sam_weights/{det_type}_detect_300e_best.pt'
sam_checkpoint = f'./xami_model/train/weights/{det_type}_sam_weights/{det_type}_sam.pth'

detr_sam_pipeline = InferXami(
    device='cuda:0',
    detr_checkpoint=detr_checkpoint,
    sam_checkpoint=sam_checkpoint,
    model_type='vit_t', # the SAM checkpoint and model_type (vit_h, vit_t, etc.) must be compatible
    use_detr_masks=True,
    detr_type=det_type)

masks = detr_sam_pipeline.run_predict('./example_images/S0893811101_M.png', show_masks=True)

For training the model, check the training README.md.

© Licence

This project is licensed under MIT license.