Sync from GitHub via hub-sync

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 innoisys
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,3 +1,377 @@
----
-license: mit
----
+# Interpretable Similarity of Synthetic Image Utility (IUS)
+[![GitHub stars](https://img.shields.io/github/stars/innoisys/ius.svg?style=flat&label=Star)](https://github.com/innoisys/ius/)
+[![Readme](https://img.shields.io/badge/README-green.svg)](README.md)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+
+
+This repository contains the official **PyTorch** implementation of the *IUS* measure, introduced in [**"Interpretable Similarity of Synthetic Image Utility"**](https://ieeexplore.ieee.org/document/11458792) , published in **IEEE Transactions on Medical Imaging (TMI)**.
+The original **TensorFlow** implementation will also be made available soon. 
+
+## Overview
+**IUS** (*Interpretable Utility Similarity*) is an interpretable measure for assessing the utility of synthetic medical image datasets for downstream clinical decision support (CDS) tasks, built upon the **EPU-CNN** framework introduced in [*E pluribus unum interpretable convolutional neural networks*](https://www.nature.com/articles/s41598-023-38459-1) and available in our [previous code repository](https://github.com/innoisys/epu-cnn-torch).
+
+## Table of Contents
+- [Features](#features)
+- [Installation](#installation)
+- [Project Structure](#project-structure)
+- [Usage](#usage)
+  - [1. Prepare Configuration](#1-prepare-configuration)
+  - [2. Supported Dataset Structures for EPU-CNN Training](#2-supported-dataset-structures-for-epu-cnn-training)
+    - [2.1 Filename-based Structure](#21-filename-based-structure)
+    - [2.2 Folder-based Structure](#22-folder-based-structure)
+    - [2.3 MedMNIST Benchmark Structure](#23-medmnist-benchmark-structure)
+  - [3. EPU-CNN Training](#3-epu-cnn-training)
+  - [4. EPU-CNN Evaluation](#4-epu-cnn-evaluation)
+  - [5. Feature contribution profile estimation](#5-feature-contribution-profile-estimation)
+  - [6. Synthetic image utility evaluation with IUS](#6-synthetic-image-utility-evaluation-with-ius)
+- [Paper Abstract](#-paper-abstract)
+- [Citation](#citation)
+- [License](#-license)
+- [TODO](#todo)
+
+## Features
+- YAML-based configuration for reproducible training, inference, and evaluation
+- Flexible support for multiple training dataset structures
+- Integrated **TensorBoard** logging for experiment monitoring
+- Complete implementation of the **IUS** pipeline:
+  1. **EPU-CNN training** and classification performance evaluation
+  2. **Baseline feature contribution profile** estimation
+  3. **Synthetic dataset or singleton evaluation using IUS** 
+- Support for both **color** and **grayscale** medical imaging modalities
+- Reproducible experiment management through structured saving of checkpoints, logs, timestamped results, and evaluation outputs
+
+
+## Installation
+Clone the repository:
+```bash
+git clone https://github.com/innoisys/ius.git
+cd ius
+```
+Install dependencies:
+```bash
+  pip install -r requirements.txt
+```
+
+## Project Structure
+```
+ius/
+├── configs/                        # YAML configuration files
+├── data/                           # Data loading and preprocessing
+│   ├── data_utils.py                   # Common utilities for image/data handling
+│   ├── dataloader.py                   # Dataloader definitions
+│   ├── dataset.py                      # Dataset implementations
+│   ├── loading.py                      # Dataset setup from YAML configuration
+│   ├── parsers.py                      # Data parser implementations
+│   └── perceptual_transforms.py        # PFM (Perceptual Feature Map) generation
+├── datasets/                       # Real datasets for EPU-CNN training & baseline feature contribution profile estimation
+├── datasets_synthetic/             # Synthetic data for IUS evaluation
+├── ius/                          # IUS implementation
+│   ├── ius.py                          # IUS measure class
+│   └── ius_eval_parser.py              # Suggested synthetic data parser (not requiring label information)
+├── model/                          # EPU-CNN model implementation
+│   ├── epu.py                          # Main EPU-CNN model definition
+│   ├── module_mapping.py               # Mappings from YAML config names to torch.nn layers/activations
+│   ├── register_modules.py             # Registry for configurable model components
+│   ├── subnetworks.py                  # Subnetwork implementation
+│   └── subnetwork_utilities.py         # Subnetwork helper modules
+├── results/                        # Training and inference outputs
+│   ├── cb_vectors.py                   # Saved baseline feature contribution profiles (from infer_cb_vector.py)
+│   ├── checkpoints.py                  # Saved EPU-CNN checkpoints and training configurations (from train_epu.py)
+│   ├── classification_performance.py   # Classification performance results (from infer_epu.py)
+│   ├── ius_eval.py                     # IUS evaluation results (from eval_ius.py)
+│   └── logs.py                         # TensorBoard logs (from train_epu.py )
+├── scripts/                        # Training and inference scripts
+│   ├── eval_ius.py                     # Runs synthetic data evaluation with IUS
+│   ├── infer_cb_vector.py              # Estimates baseline feature contribution profiles
+│   ├── infer_epu.py                    # Runs EPU-CNN inference/evaluation 
+│   └── train_epu.py                    # Trains EPU-CNN models
+└── utils/                          # Utility functions
+    ├── callbacks.py                    # Training callbacks
+    ├── config_utils.py                 # YAML/configuration utilities
+    ├── early_stopping.py               # Early stopping logic
+    ├── eval_utils.py                   # Utilities for EPU-CNN evaluation scripts
+    ├── metrics.py                      # Classification performance metrics 
+    ├── omega_parser.py                 # OmegaConf-based configuration parser 
+    ├── sanity_utils.py                 # Configuration validation and sanity checks
+    ├── tensorboard.py                  # Tensorboard utilities
+    ├── train_utils.py                  # Training setup and helping utilities
+    └── trainer.py                      # Main training loop implementation
+```
+
+## Usage
+### 1. Prepare Configuration
+
+Create a YAML configuration file in `configs/` with the following structure:
+
+```yaml
+model:
+    num_subnetworks: 4                      # set to 4, corresponds to number of perceptual feature maps, 
+    num_classes:     1
+    epu_activation: "sigmoid"               
+    subnetwork_config:
+        architecture:   "base_one"          # default ius backbone
+        input_channels: 1                   # number of channels in perceptual feature decomposition, set to 1
+        base_channels:  32
+        fc_hidden_units: 64
+        pred_activation: "tanh"
+data_params:
+    dataset_path: "../datasets/dataset_name"
+    images_extension: "jpg"
+    data_loading:
+        batch_size:         64
+        shuffle:            true
+        num_workers:        0
+        pin_memory:         false
+        persistent_workers: false
+    data_preprocessing:
+        data_mode: "rgb"                        # "rgb" or "grayscale"
+        data_parser: "filename"                 # "filename" or "folder" or "medmnist"
+        resize_dims: [128, 128]
+        medmnist_csv_file:  None
+        label_mapping:
+            abnormal: 1
+            normal: 0
+train_params:
+    mode:       "binary"
+    loss:       "binary_cross_entropy"
+    epochs:     200
+    optimizer:  "sgd"
+    learning_rate:  0.001 
+    momentum:       0.9
+    weight_decay:   0.001
+    early_stopping_patience: 30
+    early_stopping_monitor:  "val_loss"         # "val_loss" or "val_metrics.auc"
+    early_stopping_mode:     "min"              # "min" or "max"
+log_dir:          "../results/logs"             # default parent path
+checkpoint_dir:   "../results/checkpoints"      # default parent path
+experiment_name:  "ius_dataset_name"            # desired experiment path
+```
+
+**Key Points:**
+- For binary classification: Use `n_classes: 1`,  `epu_activation: "sigmoid"`
+- Adjust `input_size` and `batch_size` based on your GPU memory
+- Set `label_mapping` according to your dataset classes
+
+### 2. Supported Dataset Structures for EPU-CNN Training
+
+EPU-CNN training supports multiple dataset organization patterns. Complete examples are provided below:
+
+#### 2.1 Filename-based Structure
+
+```
+datasets/
+├── dataset_name/
+    ├── train/
+    │   ├── abnormal_001.jpg        
+    │   ├── abnormal_002.jpg        
+    │   ├── normal_001.jpg
+    │   ├── normal_002.jpg
+    ├── validation/
+    │   ├── abnormal_003.jpg        
+    │   ├── normal_003.jpg
+    └── test/
+        ├── abnormal_004.jpg        
+        └── normal_004.jpg
+```
+
+**Configuration Example:**
+```yaml
+data_params:
+    dataset_path: "../datasets/dataset_name"
+    images_extension: "jpg"
+    data_preprocessing:
+        data_mode: "rgb"                        # "rgb" or "grayscale"
+        data_parser: "filename"                 
+        resize_dims: [128, 128]
+        medmnist_csv_file:  None
+        label_mapping:
+            abnormal: 1
+            normal: 0
+```
+
+**Key Points:**
+- Requires a consistent organization across train/validation/test splits
+- The validation folder name must be exactly "validation" (not "val" or other variants)
+- Supported image formats include jpg, jpeg, png, etc, supported by the parser
+- Class names defined in label_mapping must appear in filenames
+
+#### 2.2 Folder-based Structure
+
+```
+datasets/
+├── dataset_name/
+    ├── train/
+    │   ├── abnormal   
+    │   │   ├── image_001.jpg
+    │   │   └── image_002.jpg   
+    ���   └── normal
+    │       ├── image_001.jpg
+    │       └── image_002.jpg
+    ├── validation/
+    │   ├── abnormal   
+    │   │   └── image_003.jpg
+    │   └── normal
+    │       └── image_003.jpg
+    └── test/
+        ├── abnormal   
+        │   └── image_004.jpg
+        └── normal
+            └── image_004.jpg
+```
+
+**Configuration Example:**
+```yaml
+data_params:
+    dataset_path: "../datasets/dataset_name"
+    images_extension: "jpg"
+    data_preprocessing:
+        data_mode: "rgb"                        # "rgb" or "grayscale"
+        data_parser: "folder"                 
+        resize_dims: [128, 128]
+        medmnist_csv_file:  None
+        label_mapping:
+            abnormal: 1
+            normal: 0
+```
+
+**Key Points:**
+- Requires a consistent organization across train/validation/test splits
+- The validation folder name must be exactly "validation" (not "val" or other variants)
+- Supported image formats include jpg, jpeg, png, etc. supported by the parser
+- Class names defined in label_mapping must match the class folder names
+
+#### 2.3 MedMNIST Benchmark Structure
+
+```
+datasets/
+├── pneumoniamnist/
+│   ├── pneumoniamnist.csv   
+│   ├── test_0_0.png   
+│   ├── test_1_1.png   
+│   ├── train_0_1.png   
+│   ├── train_1_0.png   
+│   ├── train_2_0.png   
+│   ├── train_3_1.png   
+│   ├── val_0_1.png   
+│   └── val_1_0.png  
+```
+
+**Configuration Example:**
+```yaml
+data_params:
+    dataset_path: "../datasets/pneumoniamnist"
+    images_extension: "png"
+    data_preprocessing:
+        data_mode: "rgb"                        # "rgb" or "grayscale"
+        data_parser: "folder"                 
+        resize_dims: [128, 128]
+        medmnist_csv_file:  "../datasets/pneumoniamnist/pneumoniamnist.csv"
+        label_mapping:
+            pneumonia: 1
+            normal: 0
+```
+
+**Key Points:**
+- Supports any 2D dataset included in the MedMNIST Benchmark collection.
+- Requires the data in the original format provided by the [MedMNIST benchmark](https://github.com/MedMNIST/MedMNIST/tree/main)
+- The CSV file should be used as downloaded, without modification.
+- Class names defined in label_mapping must match the official class names for the corresponding dataset provided in [info](https://github.com/MedMNIST/MedMNIST/blob/main/medmnist/info.py)
+
+### 3. EPU-CNN Training
+To train EPU-CNN, run one of the following commands after setting up a `config.yaml` file:
+```bash
+
+# Basic training
+python scripts/train_epu.py --config_filepath configs/train_config.yaml
+
+# Training with TensorBoard monitoring
+python scripts/train_epu.py --config_filepath configs/train_config.yaml --tensorboard
+```
+The script saves trained model checkpoints and YAML training config in results/checkpoints under automatically generated name as:
+{experiment_name}_{subnetwork_backbone}_{run_id}_{timestamp}
+
+When using the `--tensorboard` flag, the script automatically:
+- launches TensorBoard as a subprocess
+- sets up monitoring for the `logs` directory
+- makes TensorBoard available at `http://localhost:6006`
+- enables real-time monitoring of training metrics and model graphs.
+
+### 4. EPU-CNN Evaluation
+To assess the classification performance of a trained EPU-CNN model run:
+```bash
+python scripts/infer_epu.py --experiment_folder_name ius_dataset_name_base_one_0000_timestamp
+```
+ - Argument --experiment_folder_name contains the saved checkpoint and saved YAML config used during the EPU-CNN training.
+ - The classification performance report is saved in JSON format under results/classification_performance/{experiment_folder_name}
+
+### 5. Feature contribution profile estimation
+For each trained EPU-CNN model instance, the baseline feature contribution profile must be estimated only once using ether:
+```bash
+# Estimates all baseline feature contribution profiles  
+python scripts/infer_cb_vector.py --experiment_folder_name ius_dataset_name_base_one_0000_timestamp
+
+# Estimates the baseline feature contribution profile of a single class
+python scripts/infer_cb_vector.py --experiment_folder_name ius_dataset_name_base_one_0000_timestamp --cb_data normal
+```
+- The --experiment_folder_name contains the saved checkpoint and saved YAML config used during the EPU-CNN training.
+- If argument --cb_data is not provided, the script estimates the baseline feature contribution profile for each class defined in saved config YAML (automatically retrieved from results/checkpoints based on --experiment_folder_name)
+
+### 6. Synthetic image utility evaluation with IUS
+To evaluate the utility of synthetic images using IUS, run:
+```bash
+# For synthetic dataset IUS evaluation
+python scripts/eval_ius.py --experiment_folder_name ius_dataset_name_base_one_0000_timestamp --cb_vector_tag normal --synthetic_images datasets_synthetic/dataset_name/normal --synthetic_img_extension png
+
+# For a single image IUS evaluation
+python scripts/eval_ius.py --experiment_folder_name ius_dataset_name_base_one_0000_timestamp --cb_vector_tag normal --synthetic_images datasets_synthetic/dataset_name/normal/seed_000.png --synthetic_img_extension png
+```
+
+Both commands produce two outputs saved under results/ius_eval/{experiment_folder_name}
+  1. A JSON report containing information about IUS evaluation.  
+  2. A CSV file containing the estimated IUS score for  each synthetic image.
+
+**Key Points:**
+- Synthetic data should be stored under datasets_synthetic
+- No specific internal folder structure is required; only a valid path to a folder of images is needed.
+- A suggested organization example is provided below 
+- Supported image formats include jpg, jpeg, png, etc, supported by the parser
+- Class labels are neither required nor used during  IUS evaluation
+
+Suggested synthetic data structure
+```
+datasets_synthetic/
+├── dataset_name/
+    ├── normal/
+        ├── seed_000.png
+        └── seed_001.png
+```
+
+## 📄 Paper Abstract
+<p align="justify">
+Synthetic medical image data can unlock the potential of deep learning (DL)-based clinical decision support (CDS) systems through the creation of large scale, privacy-preserving, training sets. Despite the significant progress in this field, there is still a largely unanswered research question: “How can we quantitatively assess the similarity of a synthetically generated set of images with a set of real images in a given application domain?”. Today, answers to this question are mainly provided via user evaluation studies, inception-based measures, and the classification performance achieved on synthetic images. This paper proposes a novel measure to assess the similarity between synthetically generated and real sets of images, in terms of their utility for the development of DL based CDS systems. Inspired by generalized neural additive models, and unlike inception-based measures, the proposed measure is interpretable (Interpretable Utility Similarity, IUS), explaining why a synthetic dataset could be more useful than another one in the context of a CDS system based on clinically relevant image features. The experimental results on publicly available benchmark datasets from various color medical imaging modalities including endoscopic, dermoscopic and fundus imaging, indicate that selecting synthetic images of high utility similarity using IUS can result in relative improvements of up to 54.6% in terms of classification performance. The generality of IUS for synthetic data assessment is demonstrated also for grayscale X-ray and ultrasound imaging modalities. IUS implementation is available at https://github.com/innoisys/ius. 
+</p>
+
+## Citation
+If you use this code or find our work useful in your research, please cite:
+
+**APA**
+
+P. Gatoula, G. Dimas and D. K. Iakovidis, "Interpretable Similarity of Synthetic Image Utility," in *IEEE Transactions on Medical Imaging*, doi: 10.1109/TMI.2026.3679527. 
+
+**BibTeX**
+```bibtex
+@article{
+  author    = {Panagiota Gatoula and George Dimas and Dimitris K. Iakovidis},
+  title     = {Interpretable Similarity of Synthetic Image Utility},
+  journal   = {IEEE Transactions on Medical Imaging},
+  year      = {2026},
+  publisher = {IEEE},
+  doi       = {10.1109/TMI.2026.3679527}
+}
+```
+
+## ⚖️  License
+This project is licensed under the MIT License.
+
+## TODO
+- [X] Add bar-plot visualization items for the feature contribution profiles 

configs/train_config.yaml

@@ -0,0 +1,44 @@
+model:
+    num_subnetworks: 4                      # corresponds to number of perceptual feature maps, set to 4
+    num_classes:     1
+    epu_activation: "sigmoid"
+    subnetwork_config:
+        architecture:   "base_one"          # default ius backbone
+        input_channels: 1                   # number of channels in perceptual feature decomposition, set to 1
+        base_channels:  64
+        fc_hidden_units: 32
+        pred_activation: "tanh"
+
+data_params:
+    dataset_path: "../datasets/kid"
+    images_extension: "png"
+    data_loading:
+        batch_size:         8
+        shuffle:            true
+        num_workers:        0
+        pin_memory:         false
+        persistent_workers: false
+    data_preprocessing:
+        data_mode: "rgb"                        # "rgb" or "grayscale"
+        data_parser: "folder"                   # "filename" or "folder" or "medmnist"
+        resize_dims: [ 128, 128 ]
+        medmnist_csv_file:  None
+        label_mapping:
+            abnormal: 1
+            normal: 0
+
+train_params:
+    mode:       "binary"
+    loss:       "binary_cross_entropy"
+    epochs:     200
+    optimizer:  "sgd"
+    learning_rate:  0.001
+    momentum:       0.9
+    weight_decay:   0.001
+    early_stopping_patience: 10
+    early_stopping_monitor:  "val_loss"         # "val_loss" or "val_metrics.auc"
+    early_stopping_mode:     "min"              # "min" of "max"
+
+log_dir:          "../results/logs"             # default parent path
+checkpoint_dir:   "../results/checkpoints"      # default parent path
+experiment_name:  "ius_kid"

data/.gitkeep

@@ -0,0 +1 @@
+

data/data_utils.py

@@ -0,0 +1,23 @@
+from PIL import Image
+from typing import Dict
+
+
+def load_image(img_path):
+    img = Image.open(img_path).convert('RGB')
+    return img
+
+
+class LabelTransform:
+    def __init__(self, mapping_dict: Dict):
+        self.mapping_dict = mapping_dict
+        self.mapping_dict = {k.lower(): v for k, v in mapping_dict.items()}
+        self._keys = sorted(self.mapping_dict.keys())
+        # self._keys = [item.lower() for item in self._keys]
+
+    def __call__(self, label):
+        label = label.lower()
+        assert label in self._keys, (f'label {label} not in label mapping_dict provided.'
+                                     f'Available keys {self._keys}')
+        return self.mapping_dict[label]
+
+        

data/dataloader.py

@@ -0,0 +1,20 @@
+from typing import Union
+from torch.utils.data import Dataset
+from torch.utils.data import DataLoader
+
+from utils.omega_parser import DataLoading
+from data.dataset import EPUDataset
+from data.loading import EPUDatasetFromConfig
+
+
+def to_dataloader(dataset: Union[Dataset, EPUDataset, EPUDatasetFromConfig],
+                  loading_cfg: DataLoading) -> DataLoader:
+
+    return DataLoader(dataset,
+                      batch_size=loading_cfg.batch_size,
+                      shuffle=loading_cfg.shuffle,
+                      num_workers=loading_cfg.num_workers,
+                      pin_memory=loading_cfg.pin_memory,
+                      persistent_workers=loading_cfg.persistent_workers,
+                      drop_last=False
+                      )

data/dataset.py

@@ -0,0 +1,38 @@
+from torch.utils.data import Dataset
+
+from .data_utils import load_image
+
+from .parsers import BaseParser
+from .data_utils import LabelTransform
+from .perceptual_transforms import PerceptualFeatureMapTransform
+
+
+class EPUDataset(Dataset):
+    def __init__(self,
+                 data_parser: BaseParser,
+                 perceptual_transform: PerceptualFeatureMapTransform = None,
+                 label_transform: LabelTransform = None,
+                 **kwargs):
+
+        self.data_parser = data_parser
+        self.transform = perceptual_transform
+        self.label_transform = label_transform
+        self.kwargs = kwargs
+
+        self.image_paths = self.data_parser.image_filenames
+        self.labels = self.data_parser.labels
+
+    def __len__(self):
+        assert len(self.image_paths) == len(self.labels), "Mismatch in image paths and labels"
+        return len(self.image_paths)
+
+    def __getitem__(self, idx):
+        img_path = self.image_paths[idx]
+        label = self.labels[idx]
+
+        img = load_image(img_path)
+        if self.transform is not None:
+            img = self.transform(img)
+        if self.label_transform is not None:
+            label = self.label_transform(label)
+        return img, label

data/loading.py

@@ -0,0 +1,86 @@
+# from torch.utils.data import Dataset
+from typing import Union
+
+from data.data_utils import LabelTransform
+from utils.omega_parser import DataParams
+from data.parsers import set_parser_class
+from data.perceptual_transforms import PerceptualFeatureMapTransform
+from data.dataset import EPUDataset
+from utils.config_utils import data_cfg_to_dataparser
+from ius.ius_eval_parser import IUSEvalParser
+
+
+# Creates an EPUDataset from IUS config.
+# Used in EPUCNN training, classification performance eval & calculation of cb vectors
+class EPUDatasetFromConfig:
+    def __init__(self, dataconfig: DataParams, **kwargs):
+
+        self.dataset_path = dataconfig.dataset_path
+        self.images_extension = dataconfig.images_extension
+
+        self.data_preprocessing = dataconfig.data_preprocessing
+        # self.data_loading = dataconfig.data_loading
+
+        self.group_by = kwargs.get('group_by')
+
+    def get_dataset(self, dataset_mode: str) -> EPUDataset:
+        assert dataset_mode in ["train", "validation", "test"], "Dataset mode must be either train or val or test."
+
+        # Create parser & transforms for Dataset
+        parser = set_parser_class(
+            name=self.data_preprocessing.data_parser)(
+            **data_cfg_to_dataparser(
+                dataset_path=self.dataset_path,
+                images_extension=self.images_extension,
+                data_mode=dataset_mode,
+                preprocessing_cfg=self.data_preprocessing,
+                group_by=self.group_by,
+            )
+        )
+
+        perceptual_transform = PerceptualFeatureMapTransform(
+            resize_dims=self.data_preprocessing.resize_dims,
+            resize_mode="bicubic",
+            data_mode=self.data_preprocessing.data_mode
+        )
+
+        label_transform = LabelTransform(
+            mapping_dict=self.data_preprocessing.label_mapping
+        )
+
+        # Create Dataset
+        dataset = EPUDataset(
+            data_parser=parser,
+            perceptual_transform=perceptual_transform,
+            label_transform=label_transform
+        )
+        return dataset
+
+
+# used during IUS evaluation
+class IUSEvalDataset:
+    def __init__(self, dataconfig: DataParams, **kwargs):
+
+        self.dataset_path = dataconfig.dataset_path
+        self.images_extension = dataconfig.images_extension
+
+        self.data_preprocessing = dataconfig.data_preprocessing
+        # self.data_loading = dataconfig.data_loading
+
+        self.group_by = kwargs.get('group_by')
+
+    def get_dataset(self, parser: Union[IUSEvalParser]) -> EPUDataset:
+
+        perceptual_transform = PerceptualFeatureMapTransform(
+            resize_dims=self.data_preprocessing.resize_dims,
+            resize_mode="bicubic",
+            data_mode=self.data_preprocessing.data_mode
+        )
+
+        # Create Dataset
+        dataset = EPUDataset(
+            data_parser=parser,
+            perceptual_transform=perceptual_transform,
+            label_transform=None
+        )
+        return dataset

data/parsers.py

@@ -0,0 +1,301 @@
+import re
+import os.path
+import pandas as pd
+
+from glob import glob
+from typing import List, Dict
+from abc import ABC, abstractmethod
+
+
+PARSER_REGISTRY = {
+
+}
+
+
+def register_parser(name: str):
+    def decorator(cls):
+        key = name.lower()
+        if key in PARSER_REGISTRY:
+            raise ValueError(f'Parser name {key} exists already')
+        PARSER_REGISTRY[key] = cls
+        return cls
+    return decorator
+
+
+class BaseParser:
+    def __init__(self,
+                 dataset_folder: str,               # e.g datasets/dataset_name/
+                 mode: str,
+                 image_ext: str = 'jpg',
+                 **kwargs
+                 ):
+
+        self._dataset_folder = dataset_folder
+        self._image_ext = image_ext
+
+        assert mode in ['train', 'validation', 'test'], "Mode must be one of 'train', 'val', 'test'"
+        self._mode = mode
+        # self._kwargs = kwargs
+
+        # For storing (x, y) items
+        self._img_filenames = []            # filenames / item x in dataloader
+        self._labels = []                   # labels    / item y in dataloader
+
+        # Methods
+        self.parse_dataset_folder()
+
+        # Validity check
+        assert len(self._img_filenames) == len(self._labels), \
+            (f"Mismatch in number of images ({len(self._img_filenames)}) and labels ({len(self._labels)}) found "
+             f"in folder {self._dataset_folder}")
+
+        self._group_by_key = kwargs.get("group_by_key", None)
+        self._group_by_value = kwargs.get("group_by_value", None)
+        self._label_mapping = kwargs.get("label_mapping", None)
+        self.group_dataset_by()
+
+    @abstractmethod
+    def parse_dataset_folder(self):
+        raise NotImplementedError("Method must be implemented in child class")
+
+    @property
+    def image_filenames(self) -> List[str]:
+        return self._img_filenames
+
+    @property
+    def labels(self) -> List[str]:
+        return self._labels
+
+    def update_image_filenames(self, filenames: List[str]) -> None:
+        self._img_filenames = filenames
+
+    def update_labels(self, labels: List[str]) -> None:
+        self._labels = labels
+
+    def group_dataset_by(self):
+        if self._group_by_key is None and self._group_by_value is None:
+            pass
+        elif self._group_by_key is not None and self._group_by_value is not None:
+            raise ValueError(
+                "Please specify either 'group_key' or 'group_value' to perform group_dataset_by, not both"
+            )
+        else:
+            if self._label_mapping is None:
+                raise ValueError(
+                    "No 'label_mapping' specified for perform group_dataset_by() "
+                )
+
+            filtered_filenames, filtered_labels = [], []
+
+            if self._group_by_key is not None:
+                if not isinstance(self._group_by_key, list):
+                    self._group_by_key = [self._group_by_key]
+                # Valid keys
+                _label_mapping_keys = self._label_mapping.keys()
+
+                # Filter (filenames & labels) if a label is in self._group_by_key and exists in label_mapping
+                for _group_key in self._group_by_key:
+                    if _group_key not in _label_mapping_keys:
+                        raise ValueError(
+                            f"{_label_mapping_keys} does not match any available label {_group_key}"
+                        )
+                # _valid_group_values = [self._label_mapping[_group_key] for _group_key in self._group_by_key]
+                for i in range(len(self.labels)):
+                    if self.labels[i] in self._group_by_key:
+                        filtered_filenames.append(self._img_filenames[i])
+                        filtered_labels.append(self.labels[i])
+
+            elif self._group_by_value is not None:
+                _label_mapping_values = self._label_mapping.values()
+                assert self._group_by_value in _label_mapping_values, \
+                    f"Label mapping dict has values {_label_mapping_values}"
+                for i in range(len(self.labels)):
+                    if int(self.labels[i]) == int(self._group_by_value):
+                        filtered_filenames.append(self._img_filenames[i])
+                        filtered_labels.append(self.labels[i])
+
+            self.update_labels(filtered_labels)
+            self.update_image_filenames(filtered_filenames)
+
+
+@register_parser('filename')
+class FilenameParser(BaseParser):
+    """
+        expected data structure:
+                                datasets/dataset_name/mode
+                                |
+                                |----label_one*.image_ext
+                                |----label_two*.image_ext
+    """
+    def __init__(self,
+                 dataset_folder: str,               # e.g datasets/dataset_name
+                 mode: str = 'train',
+                 image_ext: str = 'jpg',
+                 **kwargs,
+                 ):
+
+        self._label_mapping = kwargs.get("label_mapping", None)
+        if self._label_mapping is None:
+            raise ValueError(
+                f"No 'label_mapping' specified for data parsing in FilenameParser. "
+            )
+
+        self._group_by_key = kwargs.get("group_by_key", None)
+        self._group_by_value = kwargs.get("group_by_value", None)
+        if self._group_by_key is None and self._group_by_value is not None:
+            raise ValueError(
+                "Use argument 'group_key' to perform group_dataset_by() in FilenameParser "
+            )
+
+        super().__init__(dataset_folder=dataset_folder,               # e.g dataset/train/
+                         mode=mode,
+                         image_ext=image_ext,
+                         **kwargs)
+
+    @staticmethod
+    def extract_labels_from_file(filepath_list: List[str], label_mapping: Dict, extension: str) -> List[str]:
+        unique_labels = sorted(label_mapping.keys())
+
+        labels = []
+        for filepath in filepath_list:
+            filename = os.path.basename(filepath).lower()
+            is_match = False
+            for label_name in unique_labels:
+                # {label}*.{ext} or {label}.{ext}
+                pattern = rf"^{re.escape(label_name)}.*\.{re.escape(extension)}$"
+                if re.match(pattern, filename):
+                    label = label_name              # label = label_mapping[label_name] moved to Label Transform
+                    labels.append(label)
+                    is_match = True
+                    break
+            if not is_match:
+                raise ValueError(
+                    f"Filename {filepath} does not match any available label {label_mapping}"
+                )
+        return labels
+
+    def parse_dataset_folder(self) -> None:
+        self._img_filenames = sorted(
+            glob(os.path.join(self._dataset_folder, self._mode, '*.{}'.format(self._image_ext)))
+        )
+        self._labels = self.extract_labels_from_file(filepath_list=self._img_filenames,
+                                                     label_mapping=self._label_mapping,
+                                                     extension=self._image_ext)
+
+
+@register_parser("folder")
+class FolderParser(BaseParser):
+    """
+        expected data structure:
+                                datasets/dataset_name/mode
+                                |----label_one
+                                     |----filename_one.image_ext
+                                     |----filename_two.image_ext
+                                |----label_two
+                                     |----filename_three.image_ext
+                                     |----filename_four.image_ext
+    """
+    def __init__(self,
+                 dataset_folder: str,               # e.g datasets/dataset_name/
+                 mode: str = 'train',
+                 image_ext: str = 'jpg',
+                 **kwargs,
+                 ):
+
+        self._label_mapping = kwargs.get("label_mapping", None)
+        if self._label_mapping is None:
+            raise ValueError(
+                f"No 'label_mapping' specified for data parsing in FolderParser. "
+            )
+
+        self._group_by_key = kwargs.get("group_by_key", None)
+        self._group_by_value = kwargs.get("group_by_value", None)
+        if self._group_by_key is None and self._group_by_value is not None:
+            raise ValueError(
+                "Use argument 'group_by_key' to perform group_dataset_by() in FolderParser "
+            )
+
+        super().__init__(dataset_folder=dataset_folder,               # e.g dataset/dataset_name/
+                         mode=mode,
+                         image_ext=image_ext,
+                         **kwargs)
+
+    def parse_dataset_folder(self) -> None:
+        unique_labels = sorted(self._label_mapping.keys())
+        unique_folders = sorted(os.listdir(os.path.join(self._dataset_folder, self._mode)))
+        # unique_folders = [item.lower() for item in unique_folders]
+        if set(unique_labels) != set(unique_folders):
+            raise ValueError(
+                f"Mismatch between label_mapping and folders in {os.path.join(self._dataset_folder, self._mode)}\n"
+                f"Folders found: {unique_folders}\n"
+                f"Labels found: {unique_labels}"
+            )
+        for category_folder in unique_folders:
+            files = glob(os.path.join(self._dataset_folder, self._mode, category_folder, '*.{}'.format(self._image_ext)))
+            self._img_filenames.extend(files)
+            labels = [category_folder] * len(files)
+            self._labels.extend(labels)
+
+
+@register_parser("medmnist")
+class MedMNISTParser(BaseParser):
+    """
+        expected data structure:
+                                datasets/dataset_name/mode
+                                |
+                                |----filename_one.image_ext
+                                |----filename_two.image_ext
+                                |----filename_three.image_ext
+                                |----filename_four.image_ext
+    """
+    def __init__(self,
+                 dataset_folder: str,               # e.g datasets/dataset_name/
+                 mode: str = 'train',
+                 image_ext: str = 'jpg',
+                 **kwargs,
+                 ):
+
+        self._csv_file = kwargs.pop("csv_file", None)
+        if self._csv_file is None:
+            raise ValueError(
+                f"No 'csv_file' specified for data parsing in CSVFileParser. "
+            )
+
+        self._group_by_key = kwargs.get("group_by_key", None)
+        self._group_by_value = kwargs.get("group_by_value", None)
+        if self._group_by_key is not None and self._group_by_value is None:
+            raise ValueError(
+                "Use argument 'group_by_value' to perform group_dataset_by() in MedMNISTParser "
+            )
+
+        super().__init__(dataset_folder=dataset_folder,               # e.g dataset/train/
+                         mode=mode,
+                         image_ext=image_ext,
+                         **kwargs)
+
+    @staticmethod
+    def read_from_medmnist_csv_file(csv_file: str, mode: str, retrieve_info: str) -> List[str]:
+        df = pd.read_csv(csv_file)
+        df.columns = ['split', 'filename', 'label']
+        df = df[df['split'] == mode.upper()]
+        return df[retrieve_info].tolist()
+
+    def parse_dataset_folder(self) -> None:
+        filenames = self.read_from_medmnist_csv_file(csv_file=self._csv_file,
+                                                     mode=self._mode,
+                                                     retrieve_info='filename')
+        filenames = [os.path.join(self._dataset_folder, file) for file in filenames]
+        self._img_filenames = filenames
+
+        self._labels = self.read_from_medmnist_csv_file(csv_file=self._csv_file,
+                                                        mode=self._mode,
+                                                        retrieve_info='label')
+
+def set_parser_class(name):
+    name = name.lower()
+    if name not in PARSER_REGISTRY:
+        raise ValueError(
+            "Unrecognized parser class name."
+            f"Please use one of the following: {PARSER_REGISTRY.keys()}"
+        )
+    return PARSER_REGISTRY[name]

data/perceptual_transforms.py

@@ -0,0 +1,164 @@
+import torch
+import numpy as np
+import torch.nn.functional as F
+
+import PIL.Image as Image
+import kornia.color as Kcolor
+import kornia.filters as Kfilters
+
+from typing import Tuple
+
+
+def to_chw_tensor(img: Image) -> torch.Tensor:
+    if not isinstance(img, torch.Tensor):
+        img = torch.tensor(np.array(img), dtype=torch.float32)
+    # img = img.float()
+
+    # Image.mode('L')
+    if img.ndim not in (2, 3):
+        raise ValueError(f'Image shape {img.shape} is not supported')
+    elif img.ndim == 2:
+        img = img.unsqueeze(0)              # [1, h, w]
+    elif img.ndim == 3:
+        img = img.permute(2, 0, 1)          # [c, h, w]
+
+    if img.shape[0] not in (1, 3):
+        raise ValueError(f'Image shape {img.shape} is not supported')
+    return img
+
+
+def resize_chw_tensor(img:torch.Tensor, resize_dims: Tuple, mode: str = "bicubic") -> torch.Tensor:
+    img = img.unsqueeze(0)                  # [1, c, h, w]
+    img = F.interpolate(img, size=resize_dims, mode=mode)
+    return img.squeeze(0)                   # [c, h, w]
+
+
+def min_max_normalize(img:torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
+    min_value = img.amin(dim=(-2, -1), keepdim=True)
+    max_value = img.amax(dim=(-2, -1), keepdim=True)
+    return (img - min_value) / (max_value - min_value + eps)
+
+
+def grayscale_perceptual_features(image: Image,
+                                  resize_dims: Tuple[int, int],
+                                  resize_mode: str = "bicubic"
+                                  ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    PFM feature extraction used in IUS measure
+    :param image:           Grayscale input image for decomposition
+    :param resize_dims:     spatial dims (h,w) for feature decomposition - output resolution
+    :param resize_mode:     method applied for resizing
+    :return:                C-F, L-D, BAND-I, BAND-II Perceptual Feature Components
+    """
+
+    image = to_chw_tensor(image)               # [c, h, w]
+    if image.shape[0] == 3:
+        image = Kcolor.rgb_to_grayscale(image.unsqueeze(0))
+        image = image.squeeze(0)               # [c, h, w]
+
+    image = resize_chw_tensor(image, resize_dims=resize_dims, mode=resize_mode)
+    image = min_max_normalize(image)
+
+    # coarse - fine
+    coarse_fine = Kfilters.sobel(image.unsqueeze(0))
+    coarse_fine = coarse_fine.squeeze(0)
+    coarse_fine = min_max_normalize(coarse_fine)
+
+    # light - dark
+    light_dark = Kfilters.gaussian_blur2d(image.unsqueeze(0),
+                                          kernel_size=(7, 7),
+                                          sigma=(3.0, 3.0),
+                                          border_type="reflect")
+    light_dark = light_dark.squeeze(0)
+    light_dark = min_max_normalize(light_dark)
+
+    # band-I & band-II
+    band_one_mask = ((image >= 0.0) & (image < 0.3)).float()
+    band_two_mask = ((image >= 0.3) & (image < 0.6)).float()
+    # band_three_mask = ((image >= 0.6) & (image < 0.9)).float()
+
+    band_one = image * band_one_mask
+    band_two = image * band_two_mask
+
+    return coarse_fine, light_dark, band_one, band_two
+
+
+def rgb_perceptual_features(image: Image,
+                            resize_dims: Tuple[int, int],
+                            resize_mode: str = "bicubic"
+                            ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    :param image:         RGB input image for decomposition
+    :param resize_mode:   method applied for resizing
+    :param resize_dims:   spatial dims (h,w) for feature decomposition - output resolution
+    :return:              R-G, B-Y, C-F, L-D Perceptual Feature Components
+    """
+
+    image = to_chw_tensor(image)                                # [c, h, w]
+    if image.shape[0] == 1:
+        image = image.repeat(3, 1, 1)
+
+    image = resize_chw_tensor(image, resize_dims=resize_dims, mode=resize_mode)
+    image = min_max_normalize(image)
+
+    # red-green & blue-yellow
+    lab_space_repr = Kcolor.rgb_to_lab(image.unsqueeze(0))       # [1, 3, h, w]
+    red_green = lab_space_repr[0, 1:2, :, :]                             # a-channel
+    blue_yellow = lab_space_repr[0, 2:3, :, :]                           # b-channel
+
+    red_green = min_max_normalize(red_green)
+    blue_yellow = min_max_normalize(blue_yellow)
+
+    # coarse - fine
+    grayscale = Kcolor.rgb_to_grayscale(image.unsqueeze(0))     # [1, c, h, w]
+    coarse_fine = Kfilters.sobel(grayscale)
+    coarse_fine = coarse_fine.squeeze(0)
+    coarse_fine = min_max_normalize(coarse_fine)
+
+    # light-dark
+    light_dark = Kfilters.gaussian_blur2d(grayscale,
+                                          kernel_size=(7, 7),
+                                          sigma=(3.0, 3.0),
+                                          border_type="reflect")
+    light_dark = light_dark.squeeze(0)
+    light_dark = min_max_normalize(light_dark)
+
+    return red_green, blue_yellow, coarse_fine, light_dark
+
+
+class PerceptualFeatureMapTransform:
+    def __init__(self,
+                 resize_dims: Tuple[int, int],
+                 resize_mode: str = "bicubic",
+                 data_mode: str = "rgb",):
+
+        data_mode = data_mode.lower()
+        assert data_mode in ["rgb", "gray"], " data_mode must be either 'rgb' or 'gray'"
+
+        self.data_mode = data_mode
+        self.resize_dims = resize_dims
+        self.resize_mode = resize_mode
+
+    def __call__(self, image: Image) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        pfms = None
+        if self.data_mode == "rgb":
+            pfms = rgb_perceptual_features(image=image,
+                                           resize_dims=self.resize_dims,
+                                           resize_mode=self.resize_mode)
+        elif self.data_mode == "gray":
+            pfms = grayscale_perceptual_features(image=image,
+                                                 resize_dims=self.resize_dims,
+                                                 resize_mode=self.resize_mode)
+
+        pfms = torch.stack(pfms, dim=0)     # [4, ch, h, w] . batched from dataloader [bs, 4, ch, h, w]
+        return pfms
+
+
+if __name__ == "__main__":
+    import PIL.Image as Image
+    image = Image.open('../datasets/kvasir-capsule/test/erosion/5e59c7fdb16c4228_32325.jpg')
+    pfms = rgb_perceptual_features(image=image, resize_dims=(128, 128), resize_mode="bicubic")
+    collage = torch.cat(pfms, dim=-1).squeeze(0).cpu().numpy()
+    collage *= 255
+    Image.fromarray(collage.astype(np.uint8)).show()
+

datasets/.gitkeep

@@ -0,0 +1 @@
+

datasets_synthetic/.gitkeep

@@ -0,0 +1 @@
+

ius/ius.py

@@ -0,0 +1,125 @@
+import os
+import re
+import torch
+import numpy as np
+import torch.nn.functional as F
+
+from glob import glob
+from typing import Optional, Union, List
+from pathlib import Path
+
+
+class IUS:
+    def __init__(self,
+                 cb_path: str,
+                 cb_tag: Optional[str] = None,
+                 device: Optional[Union[str, torch.device]] = None,):
+
+        assert cb_path.endswith('.npy') or Path(cb_path).is_dir(), \
+            'cb_path should be an .npy file or a directory with .npy files '
+
+        # cb_path is an .npy file
+        if cb_path.endswith('.npy') and not Path(cb_path).is_dir():
+            self.cb_vector = np.load(cb_path)
+
+            if cb_tag is None:
+                # extract from .npy filename
+                match = re.match(r"cb_vector_(.*)\.npy", os.path.basename(cb_path))
+                assert match is not None, f"Could not extract cb_tag from filename: {cb_path}"
+                cb_tag = match.group(1)
+
+            self.cb_tag = cb_tag
+            self.cb_mapping = {str(self.cb_tag): self.cb_vector}
+            self.cb_has_value = True
+            print("cb_vector loaded !")
+        else:
+            # cb_path is a dir with .npy files
+            self._load_cb_vecs_from_path(cb_path)
+
+        self.device = device if device is not None else torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    def _load_cb_vecs_from_path(self, cb_path: str):
+        cb_files = glob(os.path.join(cb_path, 'cb_vector_*.npy'))
+        self.cb_mapping = {}
+
+        for cb_file in cb_files:
+            cb_vector = np.load(cb_file)
+            match = re.match(r"cb_vector_(.*)\.npy", os.path.basename(cb_file))
+            if match is None:
+                continue
+            tag = match.group(1)
+            self.cb_mapping[tag] = cb_vector
+
+        self.cb_vector = None
+        self.cb_tag = None
+        self.cb_has_value = False
+
+    def update_cb_vector(self, cb_vector: Union[torch.Tensor, np.ndarray]):
+        self.cb_vector = cb_vector
+
+    def update_cb_tag(self, cb_tag: str):
+        self.cb_tag = cb_tag
+
+    def update_cb_mapping(self, cb_vector: Union[torch.Tensor, np.ndarray], cb_tag: str):
+        self.cb_mapping[cb_tag] = cb_vector
+
+    def update_device(self, device: Union[str, torch.device]):
+        self.device = device
+
+    def update(self, cb_vector: Union[torch.Tensor, np.ndarray], cb_tag: str):
+        self.update_cb_vector(cb_vector)
+        self.update_cb_tag(cb_tag)
+        self.cb_has_value = True
+        
+        # add cb_vector to cb_mapping if is a new one
+        existing_keys = list(self.cb_mapping.keys())
+        if cb_tag not in existing_keys:
+            self.update_cb_mapping(cb_vector, cb_tag)
+
+    @staticmethod
+    def _ius_score(cb_vector: torch.Tensor, c_vector: torch.Tensor) -> float:
+        ius_score = F.cosine_similarity(cb_vector, c_vector, dim=0)  # cb_vec/c_vec of shape [num_subnets, fc_pred_units]
+        ius_score = ius_score.item()
+        return float(ius_score)
+
+    @staticmethod
+    def calculate_ius_across_multiple_c(cb_vector: Union[torch.Tensor, np.ndarray],
+                                        c_vectors: torch.Tensor,
+                                        device: Union[str, torch.device]) -> List[float]:
+        ius_scores = []
+
+        if not isinstance(cb_vector, torch.Tensor):
+            cb_vector = torch.as_tensor(cb_vector)      # [num_subnets, fc_pred_units]
+        # cb_vector = cb_vector.to(c_vectors.device)
+        cb_vector = cb_vector.to(device, dtype=c_vectors.dtype)
+
+        num_of_c_vecs = c_vectors.shape[0]              # [fake_images, num_subnets, fc_pred_units]
+        for i in range(num_of_c_vecs):
+            c_hat_vec = c_vectors[i]
+            c_hat_vec = c_hat_vec.to(device)
+            ius_score = IUS._ius_score(cb_vector, c_hat_vec)
+            ius_scores.append(ius_score)
+        return ius_scores
+
+    def ius_measure(self, c_vectors: torch.Tensor, cb_tag: Optional[str] = None) -> List[float]:
+        #  IUS init from cb_path that is a dir with multiple .npy & cb_tag is still None
+        if not self.cb_has_value:
+            assert cb_tag is not None, 'Please provide cb_tag'
+
+            existing_tags = self.cb_mapping.keys()
+            assert cb_tag in existing_tags, f'cb_tag provided is not available. Available cb_tags: {existing_tags}'
+
+            self.update_cb_tag(cb_tag)
+            cb_vector = self.cb_mapping[self.cb_tag]
+            self.update_cb_vector(cb_vector)
+            self.cb_has_value = True
+
+        # IUS init from .npy cb_path
+        print(f"Calculating IUS score for cb_tag = {self.cb_tag}")
+        ius_scores = self.calculate_ius_across_multiple_c(cb_vector=self.cb_vector,
+                                                          c_vectors=c_vectors,
+                                                          device=self.device)
+        return ius_scores
+
+
+

ius/ius_eval_parser.py

@@ -0,0 +1,49 @@
+import os
+import torch
+import numpy as np
+import torch.nn.functional as F
+
+from glob import glob
+from pathlib import Path
+from typing import List, Union
+
+from data.parsers import BaseParser
+
+
+VALID_EXT = [".png", ".jpg", ".jpeg", ".tiff", ".tif",]
+
+
+class IUSEvalParser(BaseParser):
+    """
+        reads files inside a folder and creates a dummy parser (files, None)
+        labels are not required for IUS evaluation.
+        in the case where path is a single filename it creates a dummy parser (filename, None)
+    """
+
+    def __init__(self,
+                 path: str,               # e.g datasets/dataset_name
+                 image_ext: str = 'jpg',
+                 **kwargs,
+                 ):
+
+        self.single_item = False
+        self.synthetic_images_path = path
+
+        if Path(path).suffix.lower() in VALID_EXT and not Path(path).is_dir():
+            self.single_item = True
+            dataset_folder = str(Path(path).parent)
+        else:
+            dataset_folder = self.synthetic_images_path
+
+        super().__init__(dataset_folder=dataset_folder,               # e.g path_to_fake_images
+                         mode='test',
+                         image_ext=image_ext,
+                         **kwargs)
+
+    def parse_dataset_folder(self) -> None:
+        if self.single_item:
+            self._img_filenames = [self.synthetic_images_path]              # [self._dataset_folder]
+        else:
+            self._img_filenames = sorted(glob(os.path.join(self.synthetic_images_path, f"*.{self._image_ext}")))
+        self._labels = [-1 for _ in range(len(self._img_filenames))]       # ignored
+

model/.gitkeep

@@ -0,0 +1 @@
+

model/epu.py

@@ -0,0 +1,147 @@
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+
+from typing import List, Tuple, Union
+
+from .subnetwork import Subnet
+from .module_mapping import layer_mapping
+from .register_modules import get_registered_model
+
+
+class BaseAdditiveNetwork(nn.Module):
+    def __init__(self,
+                 subnetworks: nn.ModuleList,
+                 num_classes: int,
+                 activation: nn.Module
+                 ):
+        super(BaseAdditiveNetwork, self).__init__()
+
+        self.subnetworks = subnetworks
+        self.n_classes = num_classes
+        self.activation = activation
+
+        self.bias = nn.Parameter(torch.randn(self.n_classes), requires_grad=True)
+
+        self.interpretations = None
+        self.output = None
+
+    def forward(self, pfms: Union[Tuple[torch.Tensor], torch.Tensor], ret_raw_logits=True) -> torch.Tensor:
+        if isinstance(pfms, torch.Tensor) and pfms.ndim == 5:
+            pfms = torch.unbind(pfms, dim=1)
+        self.interpretations = [subnetwork(_x) for _x, subnetwork in zip(pfms, self.subnetworks)]
+        output = torch.sum(torch.stack(self.interpretations), dim=0) + self.bias
+        if not ret_raw_logits:                          # in case that ret_raw_logits = False
+            output = self.activation(output)            # [bs, num_classes]
+        self.output = output
+        return output
+
+    def get_interpretations(self) -> List[torch.Tensor]:
+        return self.interpretations
+
+    def get_outputs(self) -> torch.Tensor:
+        return self.output
+
+    def get_bias(self) -> torch.Tensor:
+        return self.bias
+
+    
+class EPUCNN(BaseAdditiveNetwork):
+    def __init__(self,
+                 num_classes: int,
+                 subnetwork_name: str,
+                 num_subnetworks: int = 4,
+                 epu_activation: str = "sigmoid",
+                 subnet_activation: str = "tanh",
+                 **kwargs,
+                 ):
+
+        subnet_cfg = kwargs.pop("subnet_cfg", None)
+        default_cfg = {
+            "input_channels": 3,
+            "base_channels": 64,
+            "fc_units": 32
+        }
+        if subnet_cfg is None:
+            subnet_cfg = default_cfg.copy()
+        else:
+            _unknown_key = set(subnet_cfg.keys()) - set(default_cfg.keys())
+            if _unknown_key:
+                raise ValueError(
+                    f"Unused subnet_cfg key: {list(_unknown_key)}"
+                    f"Expected keys: {list(default_cfg.keys())}"
+                )
+            subnet_cfg = {**default_cfg, **subnet_cfg}
+
+        subnet_class = get_registered_model(subnetwork_name)
+        subnets = nn.ModuleList(
+            [subnet_class(
+                input_channels=subnet_cfg["input_channels"],
+                base_channels=subnet_cfg["base_channels"],
+                fc_hidden_units=subnet_cfg["fc_units"],
+                fc_pred_units=num_classes,
+                pred_activation=subnet_activation,
+            ) for _ in range(num_subnetworks)]
+        )
+        epu_activation = layer_mapping(epu_activation)()
+
+        super(EPUCNN, self).__init__(
+            subnetworks=subnets,
+            num_classes=num_classes,
+            activation=epu_activation,
+        )
+
+        self.subnetworks = subnets
+        self.num_classes = num_classes
+        self.subnet_name = subnetwork_name
+        self.epu_activation = epu_activation
+        self.num_subnetworks = num_subnetworks
+        self.subnet_activation = subnet_activation
+
+    def create_baseline_feature_contribution_profile(self, data_loader: DataLoader, device: str) -> torch.Tensor:
+        # [dataloader_items_unbatched, num_subnets, num_classes]
+        profiles = self.calculate_feature_contribution_profiles(data_loader=data_loader, device=device)
+        return torch.mean(profiles, dim=0)                           # [num_subnets, num_classes]
+
+    def calculate_feature_contribution_profiles(self, data_loader: DataLoader, device: str) -> torch.Tensor:
+        self.eval()
+        profiles = []
+        with torch.no_grad():
+            for x, _ in data_loader:
+                x = x.to(device)
+                x_profile = self.feature_contribution_profile(x)    # [batch, num_subnets, num_classes]
+                profiles.append(x_profile)
+        profiles = torch.cat(profiles, dim=0)                   # [dataloader_items_unbatched, num_subnets, num_classes]
+        return profiles
+
+    def feature_contribution_profile(self, pfm_tuple: Tuple[torch.Tensor]) -> torch.Tensor:
+        self.eval()
+        with torch.no_grad():
+            _ = self.forward(pfm_tuple)
+            profile = self.get_interpretations()        # List of 4 (num_subnets) items each [batch, subnet_pred_units]
+            # profile = torch.stack(profile, dim=0)              # [num_subnets, batch, num_classes]
+            # profile = torch.permute(profile, (1, 0, 2))        # [batch, num_subnets, num_classes]
+            profile = torch.stack(profile, dim=1)                # [batch, num_subnets, num_classes]
+            return profile
+
+
+if __name__ == '__main__':
+    cfg = {
+        "input_channels": 3,
+        "base_channels": 64,
+        "fc_units": 32
+    }
+    epu = EPUCNN(
+        num_classes=1,
+        subnetwork_name="base_one",
+        num_subnetworks=4,
+        subnet_activation="sigmoid",
+        epu_activation="tanh",
+        subnet_cfg=cfg
+    )
+    x = torch.randn(16, 4, 3, 32, 32)                   # [batch, num_subnets, ch, h, w]
+    y = epu(x)
+    interpretation_vec = epu.get_interpretations()
+    bias = epu.get_bias()
+
+    vector_c = epu.feature_contribution_profile(x)

model/module_mapping.py

@@ -0,0 +1,31 @@
+import torch.nn as nn
+
+
+def layer_mapping(layer: str) -> nn.Module:
+    layer = layer.lower()
+
+    mappings = {
+        # Normalization
+        "batch": nn.BatchNorm2d,
+        "instance": nn.InstanceNorm2d,
+        "layer": lambda c: nn.GroupNorm(1, c),  # nn.LayerNorm,
+
+        # Identity
+        "none": nn.Identity,
+        "linear": nn.Identity,
+
+        # Activations
+        "relu": nn.ReLU,
+        "tanh": nn.Tanh,
+        "sigmoid": nn.Sigmoid,
+        "softmax": lambda: nn.Softmax(dim=1),
+    }
+    try:
+        return mappings[layer]
+    except KeyError as e:
+        available = list(mappings.keys())
+        raise ValueError(
+            f"{layer} not found in the existing mapping."
+            f"Existing available: {available}"
+            f"Update model.model_utils.layer_mapping()"
+        ) from e

model/register_modules.py

@@ -0,0 +1,25 @@
+import torch.nn as nn
+
+REGISTERED_MODELS = {}
+
+
+def register_model(name: str):
+    def decorator(model_class: nn.Module) -> nn.Module:
+        key = name.lower()
+        if key in REGISTERED_MODELS:
+            raise ValueError(
+                f'Model {name} already registered'
+            )
+        REGISTERED_MODELS[key] = model_class
+        return model_class
+    return decorator
+
+
+def get_registered_model(name: str) -> nn.Module:
+    key = name.lower()
+    if key not in REGISTERED_MODELS:
+        raise ValueError(
+            f'Unknown model name: {name}. '
+            f'Available models: {list(REGISTERED_MODELS.keys())}'
+        )
+    return REGISTERED_MODELS[key]

model/subnetwork.py

@@ -0,0 +1,86 @@
+import torch
+import torch.nn as nn
+
+from .subnetwork_utils import BaseBlockConvBN, TopHead
+from .register_modules import register_model
+
+
+class BaseSubNetwork(nn.Module):
+    def __init__(self,
+                 input_channels: int,
+                 base_channels: int,
+                 fc_hidden_units: int,
+                 fc_pred_units: int,
+                 pred_activation: str,
+                 ):
+        super(BaseSubNetwork, self).__init__()
+
+        self.input_channels = input_channels
+        self.base_channels = base_channels
+        self.fc_hidden_units = fc_hidden_units
+        self.fc_pred_units = fc_pred_units
+        self.pred_activation = pred_activation
+
+        self.intermediate_features = None
+
+    def get_intermediate_features(self) -> torch.Tensor:
+        return self.intermediate_features
+
+
+@register_model("base_one")
+class Subnet(BaseSubNetwork):
+    def __init__(self, input_channels=3, base_channels=32, fc_hidden_units=64, fc_pred_units=1, pred_activation="sigmoid"):
+
+        super(Subnet, self).__init__(
+            input_channels=input_channels,
+            base_channels=base_channels,
+            fc_hidden_units=fc_hidden_units,
+            fc_pred_units=fc_pred_units,
+            pred_activation=pred_activation,
+        )
+
+        self.block_one = BaseBlockConvBN(in_ch=input_channels,
+                                         out_ch=base_channels,
+                                         conv_layers=2,
+                                         kernel_size=(3, 3),
+                                         stride=(2, 2),
+                                         padding=(1, 1),
+                                         activation="relu",
+                                         normalization=True,)
+
+        self.block_two = BaseBlockConvBN(in_ch=base_channels,
+                                         out_ch=base_channels*2,
+                                         conv_layers=2,
+                                         kernel_size=(3, 3),
+                                         stride=(2, 2),
+                                         padding=(1, 1),
+                                         activation="relu",
+                                         normalization=True,)
+        
+        self.block_three = BaseBlockConvBN(in_ch=base_channels*2,
+                                           out_ch=base_channels*4,
+                                           conv_layers=3,
+                                           kernel_size=(3, 3),
+                                           stride=(2, 2),
+                                           padding=(1, 1),
+                                           activation="relu",
+                                           normalization=True,)
+        self.flatten = nn.Flatten()
+
+        self.head = TopHead(fc_units=fc_hidden_units,
+                            num_classes=fc_pred_units,
+                            hidden_layers=1,
+                            dropout_rate=0.6,
+                            fc_activation="relu",
+                            pred_activation=pred_activation)
+
+        self.intermediate_features = None
+
+    def forward(self, x):
+        x = self.block_one(x)
+        x = self.block_two(x)
+        self.intermediate_features = self.block_two.get_block_feats()
+        x = self.block_three(x)
+        x = self.flatten(x)
+        x = self.head(x)
+        return x

model/subnetwork_utils.py

@@ -0,0 +1,111 @@
+import torch
+import torch.nn as nn
+
+from typing import Tuple, Optional
+
+from model.module_mapping import layer_mapping
+
+
+class BaseBlockConvBN(nn.Module):
+    def __init__(self,
+                 in_ch: int,
+                 out_ch: int,
+                 kernel_size: Tuple[int, int] = (3, 3),
+                 stride: Tuple[int, int] = (2, 2),
+                 padding: Tuple[int, int] = (1, 1),
+                 conv_layers: int = 2,
+                 pool_layer: bool = True,
+                 normalization: bool = True,
+                 activation: Optional[str] = "linear",
+                 # return_feats: bool = False,
+                 ):
+        super(BaseBlockConvBN, self).__init__()
+
+        self.conv_layers = torch.nn.ModuleList()
+        # Same params for all convolutional layers in ModuleList
+        #      activation, out_channels, kh, kw, stride, padding
+        for _ in range(conv_layers):
+            self.conv_layers.append(
+                nn.Conv2d(in_channels=in_ch,
+                          out_channels=out_ch,
+                          kernel_size=kernel_size,
+                          stride=(1, 1),
+                          padding=padding,
+                          bias=True)
+            )
+            in_ch = out_ch
+        self._activation = layer_mapping(activation.lower())()
+
+        self._batch_norm = nn.BatchNorm2d(out_ch) if normalization else None
+        self._max_pooling = nn.MaxPool2d(kernel_size=(2, 2), stride=stride) if pool_layer else None
+
+        self._intermediate_feats = None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for conv in self.conv_layers:
+            x = conv(x)
+            x = self._activation(x)
+        self._intermediate_feats = x
+
+        if self._max_pooling is not None:
+            x = self._max_pooling(x)
+
+        if self._batch_norm is not None:
+            x = self._batch_norm(x)
+
+        return x
+
+    def get_block_feats(self) -> torch.Tensor:
+        return self._intermediate_feats
+
+
+class TopHead(nn.Module):
+    def __init__(self,
+                 # in_feats: int,
+                 fc_units: int = 64,
+                 num_classes: int = 1,
+                 hidden_layers: int = 1,
+                 fc_activation: Optional[str] = "relu",
+                 pred_activation: Optional[str] = "sigmoid",
+                 dropout_rate: Optional[float] = None,
+                 ):
+        super(TopHead, self).__init__()
+
+        self._dense_layers = nn.ModuleList()
+
+        for i in range(hidden_layers + 1):
+            if i == 0:
+                self._dense_layers.append(
+                    nn.LazyLinear(out_features=fc_units, bias=True)
+                )
+            else:
+                self._dense_layers.append(
+                    nn.Linear(in_features=fc_units, out_features=fc_units, bias=True)
+                )
+            self._dense_layers.append(
+                layer_mapping(fc_activation.lower())()
+            )
+            if dropout_rate is not None:
+                self._dense_layers.append(
+                    nn.Dropout(dropout_rate)
+                )
+
+        # prediction layer
+        self._dense_layers.append(
+            nn.Linear(in_features=fc_units, out_features=num_classes, bias=True),
+        )
+        self._dense_layers.append(
+            layer_mapping(pred_activation.lower())()
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for layer in self._dense_layers:
+            x = layer(x)
+        return x
+
+
+if __name__ == "__main__":
+    block = BaseBlockConvBN(in_ch=3, out_ch=64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), normalization=True,
+                            activation="relu", conv_layers=2)
+    dummy_input = torch.randn((1, 3, 32, 32))
+    out = block(dummy_input)

requirements.txt

@@ -0,0 +1,11 @@
+kornia==0.8.2
+numpy==2.2.0
+omegaconf==2.3.0
+pandas==2.2.3
+Pillow==11.0.0
+scikit-learn==1.7.0
+tensorboard==2.19.0
+torch==2.6.0
+torchvision==0.21.0
+tqdm==4.67.1
+

scripts/.gitkeep

@@ -0,0 +1 @@
+

scripts/eval_ius.py

@@ -0,0 +1,117 @@
+import os
+import torch
+import argparse
+import numpy as np
+import pandas as pd
+
+from pathlib import Path
+from datetime import datetime
+
+# mine
+from ius.ius import IUS
+from ius.ius_eval_parser import IUSEvalParser
+from data.loading import IUSEvalDataset
+from data.dataloader import to_dataloader
+from utils.eval_utils import EPUCNNEval
+from utils.train_utils import create_output_folders, save_to_json
+from utils.omega_parser import IUSConfig
+
+
+BASE_PATH = Path(__file__).resolve().parent
+
+
+def parse_options():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--experiment_folder_name", type=str, required=True,
+                        help="Folder name containing epu configuration & saved ckpt")
+    parser.add_argument("--cb_vector_tag", type=str, default=None, required=True,
+                        help="cb vector data. If not specified estimates all cb_vectors")
+    parser.add_argument("--synthetic_images", type=str, default='png',
+                        help="It can be either a single synthetic image file or a folder path containing multiple "
+                             "synthetic images")
+    parser.add_argument("--synthetic_img_extension", type=str, default='png',
+                        help="Extension of synthetic images")
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_options()
+    cb_vec_tag = args.cb_vector_tag
+    synthetic_img_path = args.synthetic_images
+    synthetic_img_extension = args.synthetic_img_extension
+    experiment_folder_name = args.experiment_folder_name
+
+    # Load saved config
+    print('Loading configuration...')
+    saved_epu_folder = (BASE_PATH / "../results/checkpoints").resolve()
+    saved_epu_folder = os.path.join(saved_epu_folder, experiment_folder_name)
+    cfg_path = os.path.join(saved_epu_folder, "epu_config.yaml")
+    cfg = IUSConfig.from_yaml(cfg_path)
+
+    # Load EPU for obtaining Contribution Feature Profiles of synthetic images
+    print('Loading trained model...')
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    trained_ckpt = os.path.join(saved_epu_folder, f"ckpt_{experiment_folder_name}.pt")
+    epu = EPUCNNEval(epu_cfg=cfg.model)
+    trained_epu = epu.load_ckpt(device=device, ckpt_path=trained_ckpt)
+
+    # Load Baseline Contribution Feature Profile (cb_vector)
+    print('IUS class...')
+    saved_cb_vec_folder = (BASE_PATH / "../results/cb_vectors").resolve()
+    saved_cb_vec_folder = os.path.join(saved_cb_vec_folder, experiment_folder_name)
+    cb_vec_path = os.path.join(saved_cb_vec_folder, f"cb_vector_{cb_vec_tag}.npy")
+    ius = IUS(cb_path=cb_vec_path, cb_tag=cb_vec_tag, device=device)
+
+    # Load Synthetic Images & Estimate their Contribution Feature Profiles (c^_vectors)
+    synthetic_img_path = str((BASE_PATH / ".." / synthetic_img_path).resolve())
+    parser = IUSEvalParser(path=synthetic_img_path, image_ext=synthetic_img_extension,)
+    synthetic_dset = IUSEvalDataset(dataconfig=cfg.data_params).get_dataset(parser=parser)
+
+    cfg.data_params.data_loading.shuffle = False         # update cfg.data_params.data_loading
+    cfg.data_params.data_loading.batch_size = 1
+    synthetic_dataloader = to_dataloader(dataset=synthetic_dset, loading_cfg=cfg.data_params.data_loading)
+
+    synthetic_c_vecs = trained_epu.calculate_feature_contribution_profiles(
+        data_loader=synthetic_dataloader,
+        device=device,
+    )
+
+    ius_score_list = ius.ius_measure(c_vectors=synthetic_c_vecs)
+    synthetic_filenames = parser.image_filenames
+
+    # Setup Saving Dir
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    output_folder = str((BASE_PATH / "../results/ius_eval").resolve())
+    output_folder = os.path.join(output_folder, experiment_folder_name)
+    create_output_folders(output_folder)        # skipped if exists
+
+    ius_score_pd = pd.DataFrame(
+        {
+            "filename": synthetic_filenames,
+            "ius_measure_score":  ius_score_list
+        }
+    )
+    file_saved_ius_scores = os.path.join(output_folder, f"ius_scores_{timestamp}.csv")
+    ius_score_pd.to_csv(file_saved_ius_scores, index=False)
+
+    save_info_dict = {
+        "epu_ckpt": trained_ckpt,
+        "config": cfg_path,
+        "experiment_id": cfg.experiment_saved_folder_name,
+        "cb_vector_path": cb_vec_path,
+        "synthetic_images_path": synthetic_img_path,
+        "synthetic_img_extension": synthetic_img_extension,
+        "synthetic_samples_num": len(synthetic_filenames),
+        "timestamp": timestamp,
+        "ius_scores": file_saved_ius_scores,
+    }
+    json_file_path = os.path.join(output_folder, f"ius_scores_info_{timestamp}.json")
+    save_to_json(save_info_dict, json_file_path)
+
+
+if __name__ == "__main__":
+    # python - m scripts.eval_ius - -experiment_folder_name ius_dataset_name_base_one_0000_timestamp - -cb_vector_tag normal - -synthetic_images datasets_synthetic/dataset_name/normal - -synthetic_img_extension png
+    # python - m scripts.eval_ius - -experiment_folder_name ius_dataset_name_base_one_0000_timestamp - -cb_vector_tag normal - -synthetic_images datasets_synthetic/dataset_name/normal/image_0001.png - -synthetic_img_extension png
+
+    main()

scripts/infer_cb_vector.py

@@ -0,0 +1,85 @@
+import os
+import torch
+import argparse
+import numpy as np
+
+from pathlib import Path
+
+# mine
+from data.loading import EPUDatasetFromConfig
+from data.dataloader import to_dataloader
+from utils.omega_parser import IUSConfig
+from utils.train_utils import create_output_folders
+from utils.eval_utils import EPUCNNEval
+
+
+
+BASE_PATH = Path(__file__).resolve().parent
+
+
+def parse_options():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--experiment_folder_name", type=str, required=True,
+                        help="Folder name containing epu configuration & saved ckpt "
+                             "eg ius_dataset_name_base_one_0000_timestamp")
+    parser.add_argument("--cb_data", type=str, default=None,
+                        help="cb vector data. If not specified estimates all cb_vectors")
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_options()
+
+    # Load saved config
+    print('Loading configuration...')
+
+    saved_epu_folder = (BASE_PATH / "../results/checkpoints").resolve()
+    saved_epu_folder = os.path.join(saved_epu_folder, args.experiment_folder_name)
+    cfg_path = os.path.join(saved_epu_folder, "epu_config.yaml")
+    cfg = IUSConfig.from_yaml(cfg_path)
+
+    model_cfg, data_params = cfg.model, cfg.data_params
+
+    # Create model
+    print('Loading trained model...')
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    trained_ckpt = os.path.join(saved_epu_folder, f"ckpt_{cfg.experiment_saved_folder_name}.pt")
+    epu = EPUCNNEval(epu_cfg=cfg.model)
+    trained_epu = epu.load_ckpt(device=device, ckpt_path=trained_ckpt)
+
+    # Setup Saving Dir
+    output_folder = str((BASE_PATH / "../results/cb_vectors").resolve())
+    output_folder = os.path.join(output_folder, args.experiment_folder_name)
+    create_output_folders(output_folder)        # skipped if exists
+
+    # Setup Dataset & Dataloader
+    print('Load data...')
+
+    if args.cb_data is not None:
+        group_by = [args.cb_data]
+    else:
+        group_by = data_params.data_preprocessing.label_mapping.keys()
+
+    for group_by_item in group_by:
+        data_params.dataset_path = str((BASE_PATH / data_params.dataset_path).resolve())
+        dset = EPUDatasetFromConfig(dataconfig=data_params, group_by=group_by_item)
+        # eval_dataset = dset.get_dataset(dataset_mode="test")
+        eval_dataset = dset.get_dataset(dataset_mode="validation")
+        dataloader_eval = to_dataloader(dataset=eval_dataset, loading_cfg=data_params.data_loading)
+
+        cb_vector = trained_epu.create_baseline_feature_contribution_profile(
+            data_loader=dataloader_eval,
+            device=device
+        )
+        cb_vector = cb_vector.detach().cpu().numpy()
+        cb_path = os.path.join(output_folder, f"cb_vector_{group_by_item}.npy")
+        np.save(cb_path, cb_vector)
+        print(f'cb vector saved at {cb_path}')
+
+    print('cb vector computation finished.')
+
+
+if __name__ == "__main__":
+    # python -m scripts.infer_cb_vector --experiment_folder_name ius_dataset_name_base_one_0000_timestamp
+    main()

scripts/infer_epu.py

@@ -0,0 +1,106 @@
+import os
+import torch
+import argparse
+
+from pathlib import Path
+from datetime import datetime
+
+# mine
+from data.loading import EPUDatasetFromConfig
+from data.dataloader import to_dataloader
+from model.module_mapping import layer_mapping
+from utils.metrics import EPUMetrics
+
+from utils.eval_utils import EPUCNNEval
+from utils.eval_utils import InferenceRunnerEPUCNN
+from utils.train_utils import create_output_folders, save_to_json
+from utils.omega_parser import IUSConfig
+
+
+BASE_PATH = Path(__file__).resolve().parent
+
+
+def parse_options():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--experiment_folder_name", type=str, required=True,
+                        help="Folder name containing epu configuration & saved ckpt")
+    parser.add_argument("--data_split", type=str, default="test", choices=["train", "test", "validation"],)
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_options()
+    data_split = args.data_split
+    group_by = None
+
+    # Load saved config
+    print('Loading configuration...')
+    saved_epu_folder = (BASE_PATH / "../results/checkpoints").resolve()
+    saved_epu_folder = os.path.join(saved_epu_folder, args.experiment_folder_name)
+    cfg_path = os.path.join(saved_epu_folder, "epu_config.yaml")
+    cfg = IUSConfig.from_yaml(cfg_path)
+
+    # Setup Dataset & Dataloader
+    print('Load data...')
+    dset = EPUDatasetFromConfig(dataconfig=cfg.data_params, group_by=group_by)
+    eval_dataset = dset.get_dataset(dataset_mode=data_split)
+
+    # update cfg.data_params.data_loading
+    cfg.data_params.data_loading.shuffle = False
+    cfg.data_params.data_loading.batch_size = 1
+    dataloader_eval = to_dataloader(dataset=eval_dataset, loading_cfg=cfg.data_params.data_loading)
+
+    # Setup Saving Dir
+    output_folder = str((BASE_PATH / "../results/classification_performance").resolve())
+    output_folder = os.path.join(output_folder, args.experiment_folder_name)
+    create_output_folders(output_folder)        # skipped if exists
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+
+    # Load model
+    print('Loading trained model...')
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    trained_ckpt = os.path.join(saved_epu_folder, f"ckpt_{cfg.experiment_saved_folder_name}.pt")
+    epu = EPUCNNEval(epu_cfg=cfg.model)
+    trained_epu = epu.load_ckpt(device=device, ckpt_path=trained_ckpt)
+
+    trained_epu = InferenceRunnerEPUCNN(epu_model=trained_epu,
+                                        device=device,
+                                        mode=cfg.train_params.mode)
+
+    # metrics = EPUMetrics(mode=cfg.train_params.mode,
+    #                      n_classes=cfg.model.num_classes,
+    #                      activation=layer_mapping(cfg.model.epu_activation)())
+    # epu_results = trained_epu.predict(dataloader=dataloader_eval,
+    #                                   raw_logits=True,                  # not epu activation
+    #                                   return_predictions=True)
+
+    epu_results = trained_epu.predict(dataloader=dataloader_eval,
+                                      raw_logits=False,                # apply epu activation
+                                      return_predictions=True)
+    confidence_level = 0.5
+    metrics = EPUMetrics(mode=cfg.train_params.mode,
+                         n_classes=cfg.model.num_classes,
+                         confidence_level=confidence_level,
+                         activation=layer_mapping("none")())
+
+    metric_scores = metrics.compute(y_true=epu_results["targets"], y_pred=epu_results["predictions"],)
+
+    save_info_dict = {
+        "epu_ckpt": trained_ckpt,
+        "config": cfg_path,
+        "data_split": data_split,
+        "batch_size": cfg.data_params.data_loading.batch_size,
+        "dataloader_samples": len(epu_results["targets"]),
+        "classification_performance": metric_scores,
+        "confidence": confidence_level,
+        "timestamp": timestamp,
+        "experiment_id": cfg.experiment_saved_folder_name,
+    }
+    json_file_path = os.path.join(output_folder, f"epu_classification_performance_{timestamp}.json")
+    save_to_json(save_info_dict, json_file_path)
+
+
+if __name__ == "__main__":
+    # python -m scripts.infer_epu --experiment_folder_name ius_dataset_name_base_one_0000_timestamp
+    main()

scripts/train_epu.py

@@ -0,0 +1,120 @@
+import os
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import argparse
+
+from pathlib import Path
+from datetime import datetime
+
+# mine
+from utils.omega_parser import IUSConfig
+from utils.train_utils import (create_output_folders, create_experiment_folder,
+                               update_experiment_metadata, save_config_to_output_folder)
+from utils.sanity_utils import SanityChecker
+from utils.callbacks import setup_callbacks
+from utils.config_utils import model_cfg_to_epucnn
+from utils.trainer import EPUTrainer
+from utils.metrics import EPUMetrics
+from model.epu import EPUCNN
+from model.module_mapping import layer_mapping
+from data.loading import EPUDatasetFromConfig
+from data.dataloader import to_dataloader
+
+
+BASE_PATH = Path(__file__).resolve().parent
+
+
+def parse_options():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config_filepath", type=str, required=True, help="Path containing configuration")
+    parser.add_argument("--tensorboard", action="store_true", help="Enable tensorboard")
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_options()
+
+    print('Loading configuration...')
+    # Load configuration .yaml
+    cfg = IUSConfig.from_yaml(args.config_filepath)
+
+    # Sanity Check User's Config
+    SanityChecker(cfg).sanity_check()
+
+    # Set User's params
+    model_cfg, train_params, data_params = cfg.model, cfg.train_params, cfg.data_params
+
+    print('Setup directories...')
+    # Setup Experiment Name & Saving Directories
+    cfg.log_dir = str((BASE_PATH / cfg.log_dir).resolve())
+    cfg.checkpoint_dir = str((BASE_PATH / cfg.checkpoint_dir).resolve())
+    create_output_folders([cfg.log_dir, cfg.checkpoint_dir])  # skipped if exists
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    experiment_name = create_experiment_folder(log_root=cfg.log_dir, model=model_cfg.subnetwork_config.architecture,
+                                               experiment=cfg.experiment_name, timestamp=timestamp)
+    logs_folder = os.path.join(cfg.log_dir, experiment_name)
+    ckpt_folder = os.path.join(cfg.checkpoint_dir, experiment_name)
+    create_output_folders([logs_folder, ckpt_folder])
+    update_experiment_metadata(cfg, experiment_name=experiment_name, timestamp=timestamp,)
+    save_config_to_output_folder(out_folder=ckpt_folder, cfg=cfg, cfg_filename="epu_config.yaml")
+
+    # Set Device
+    print('Set device ...')
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Create model
+    print('Build detection model...')
+    epu_model = EPUCNN(**model_cfg_to_epucnn(model_cfg))
+
+    # # Setup Dataset & Dataloader
+    print('Load data...')
+    data_params.dataset_path = str((BASE_PATH / data_params.dataset_path).resolve())
+    dset = EPUDatasetFromConfig(dataconfig=data_params)
+    dataset_train = dset.get_dataset(dataset_mode="train")
+    dataset_val = dset.get_dataset(dataset_mode="validation")
+    dataloader_train = to_dataloader(dataset=dataset_train, loading_cfg=data_params.data_loading)
+    dataloader_val = to_dataloader(dataset=dataset_val, loading_cfg=data_params.data_loading)
+
+    print('Setup optimizer and callbacks ...')
+    # Setup callbacks loss & optimizer & metrics
+    calls = setup_callbacks(ckpt_path=os.path.join(ckpt_folder, f"ckpt_{experiment_name}.pt"),
+                            log_dir=logs_folder,
+                            early_patience=train_params.early_stopping_patience,
+                            early_mode=train_params.early_stopping_mode,
+                            early_monitor=train_params.early_stopping_monitor,
+                            use_tensorboard=args.tensorboard,
+                            )   # other kwargs to pass, override defaults:
+                                # delta=0, verbose=True, restore_best_weights=False,save_final_model=True)
+                                # log_histograms=False, tb_port=6006, tb_browser=False
+
+    loss_fun = nn.BCEWithLogitsLoss() if train_params.mode == "binary" else nn.CrossEntropyLoss()
+    optimizer = optim.SGD(params=epu_model.parameters(),
+                          lr=train_params.learning_rate,
+                          momentum=train_params.momentum,
+                          weight_decay=train_params.weight_decay,)
+    metrics = EPUMetrics(mode=train_params.mode,
+                         n_classes=model_cfg.num_classes,
+                         activation=layer_mapping(model_cfg.epu_activation)()
+                         )
+
+    # # launch training
+    print('Start training...')
+    trainer = EPUTrainer(model=epu_model,
+                         device=device,
+                         optimizer=optimizer,
+                         criterion=loss_fun,
+                         epochs=train_params.epochs,
+                         train_loader=dataloader_train,
+                         val_loader=dataloader_val,
+                         callbacks=calls,
+                         metrics=metrics,
+                         checkpoint_dir=ckpt_folder,
+                         )
+    trainer.train()
+
+
+if __name__ == "__main__":
+    # python -m scripts.train_epu --config_filepath configs/train_config.yaml
+    main()

utils/.gitkeep

@@ -0,0 +1 @@
+

utils/callbacks.py

@@ -0,0 +1,41 @@
+from typing import List
+
+from utils.early_stopping import EarlyStoppingCallback
+from utils.tensorboard import TensorboardLoggerCallback
+
+
+def setup_callbacks(ckpt_path: str,
+                    log_dir: str,
+                    early_patience: int,
+                    early_monitor: str = "val_loss",
+                    early_mode: str = "min",
+                    use_tensorboard: bool = False,
+                    **kwargs) -> List[object]:
+
+    delta = kwargs.get("delta", 0)
+    verbose = kwargs.get("verbose", True)
+    restore_best_weights = kwargs.get("restore_best_weights", False)
+    save_final_model = kwargs.get("save_final_model", True)
+
+    es_call = EarlyStoppingCallback(patience=early_patience,
+                                    delta=delta,
+                                    checkpoint_path=ckpt_path,
+                                    verbose=verbose,
+                                    restore_best_weights=restore_best_weights,
+                                    monitor=early_monitor,
+                                    mode=early_mode,
+                                    save_final_model=save_final_model
+                                    )
+
+    log_histograms = kwargs.get("log_histograms", False)
+    log_images_every = kwargs.get("log_images_every", 10)
+    tb_port = kwargs.get("tb_port", 6006)
+    tb_browser = kwargs.get("tb_browser", False)
+
+    tb_logger = TensorboardLoggerCallback(log_dir=log_dir,
+                                          log_histograms=log_histograms,
+                                          launch_tb=use_tensorboard,
+                                          tb_port=tb_port,
+                                          open_tb_in_browser=tb_browser
+                                          )
+    return [es_call, tb_logger]

utils/config_utils.py

@@ -0,0 +1,54 @@
+from utils.omega_parser import EPUCNNParams
+from utils.omega_parser import DataParams, DataPreprocessing
+
+
+def model_cfg_to_epucnn(model_cfg: EPUCNNParams):
+    return {
+        "num_classes": model_cfg.num_classes,
+        "subnetwork_name": model_cfg.subnetwork_config.architecture,
+        "num_subnetworks": model_cfg.num_subnetworks,
+        "epu_activation":  model_cfg.epu_activation,
+        "subnet_activation": model_cfg.subnetwork_config.pred_activation,
+        "subnet_cfg": {
+            "input_channels": model_cfg.subnetwork_config.input_channels,
+            "base_channels": model_cfg.subnetwork_config.base_channels,
+            "fc_units": model_cfg.subnetwork_config.fc_hidden_units,
+        }
+    }
+
+
+def data_cfg_to_dataparser(dataset_path: str,
+                           images_extension: str,
+                           data_mode: str,
+                           preprocessing_cfg: DataPreprocessing,
+                           group_by: str = None,):
+
+    group_by_key = None
+    group_by_value = None
+    if group_by is not None:
+        if preprocessing_cfg.data_parser in ["filename", "folder"]:
+            group_by_key = group_by                                             # eg "normal"
+        elif preprocessing_cfg.data_parser in ["medmnist"]:
+            group_by_value = preprocessing_cfg.label_mapping.get(group_by)      # eg "0" from item ("normal": "0")
+
+    return {
+        "dataset_folder": dataset_path,
+        "mode": data_mode,
+        "image_ext": images_extension,
+        "label_mapping": preprocessing_cfg.label_mapping,
+        "csv_file": preprocessing_cfg.medmnist_csv_file,
+        "group_by_key": group_by_key,
+        "group_by_value": group_by_value
+    }
+
+
+# def data_cfg_to_dataparser(data_cfg: DataParams, data_mode: str = 'train'):
+#     return {
+#         "dataset_folder": data_cfg.dataset_path,
+#         "mode": data_mode,
+#         "image_ext": data_cfg.images_extension,
+#         "kwargs": {
+#             "label_mapping": data_cfg.data_preprocessing.label_mapping,
+#             "csv_file": data_cfg.data_preprocessing.medmnist_csv_file
+#         }
+#     }

utils/early_stopping.py

@@ -0,0 +1,145 @@
+import torch
+import torch.nn as nn
+
+
+class EarlyStopping(object):
+    def __init__(self, patience: int = 10, delta: float = 0, checkpoint_path: str = 'checkpoint.pt',
+                 verbose: bool = True, restore_best_weights: bool = True, monitor: str = 'val_loss',
+                 mode: str = 'min'):
+        """
+        Args:
+            patience (int): How many epochs to wait after last improvement before stopping.
+            delta (float):  Minimum change to qualify as an improvement.
+            checkpoint_path (str): File path to save the best model.
+            verbose (bool): If True, prints messages when improvement occurs.
+            restore_best_weights (bool): If True, loads the best weights saved during training.
+            monitor (str): Metric name to monitor (e.g., 'val_loss', 'val_f1', 'val_accuracy').
+            mode (str): 'min' if lower is better (e.g. loss), 'max' if higher is better (e.g. accuracy, f1).
+        """
+
+        self.patience = patience
+        self.delta = delta
+        self.checkpoint_path = checkpoint_path
+        self.verbose = verbose
+        self.restore_best_weights_flag = restore_best_weights
+        self.monitor = monitor
+        self.mode = mode
+
+        if self.mode not in ['min', 'max']:
+            raise ValueError("mode must be 'min' or 'max'")
+
+        self.counter = 0
+        self.best_score = None
+        self.best_epoch = 0
+        self.early_stop = False
+        self.best_value = float('inf') if mode == 'min' else -float('inf')
+
+    def __call__(self, current_value: float, model: nn.Module, epoch: int = None):
+        improvement = (
+                (self.mode == 'min' and current_value < self.best_value - self.delta) or
+                (self.mode == 'max' and current_value > self.best_value + self.delta)
+        )
+
+        if self.best_score is None:
+            self.best_score = current_value
+            self.best_epoch = epoch if epoch is not None else 0
+            self.save_checkpoint(current_value, model)
+            return False
+
+        if improvement:
+            if self.verbose:
+                print(f"EarlyStopping (Improved): {self.monitor} {self.best_value:.4f} → {current_value:.4f}")
+            self.best_score = current_value
+            self.best_epoch = epoch if epoch is not None else 0
+            self.save_checkpoint(current_value, model)
+            self.counter = 0
+        else:
+            self.counter += 1
+            if self.verbose:
+                print(f"EarlyStopping (No Improvement): {self.counter}/{self.patience} epochs "
+                      f"({self.monitor}={current_value:.4f})")
+            if self.counter >= self.patience:
+                self.early_stop = True
+
+        return self.early_stop
+
+    def save_checkpoint(self, current_value: float, model: nn.Module):
+        if self.verbose:
+            print(f"Saving new best model at {self.checkpoint_path} ({self.monitor}={current_value:.4f})")
+        torch.save(model.state_dict(), self.checkpoint_path)
+        self.best_value = current_value
+
+    def restore_best_weights(self, model: nn.Module):
+        if self.restore_best_weights_flag:
+            if torch.cuda.is_available():
+                map_location = torch.device('cuda')
+            else:
+                map_location = torch.device('cpu')
+
+            try:
+                model.load_state_dict(torch.load(self.checkpoint_path, map_location=map_location))
+                if self.verbose:
+                    print(f'Restored best model from {self.checkpoint_path} (epoch {self.best_epoch})')
+            except FileNotFoundError:
+                if self.verbose:
+                    print(f'Restoring failed | No checkpoint found at {self.checkpoint_path}, cannot restore weights.')
+
+
+class EarlyStoppingCallback(object):
+    """Adapter to use EarlyStopping as a callback in the trainer function."""
+
+    def __init__(self, patience=10, delta=0, checkpoint_path='checkpoint.pt', verbose=True, restore_best_weights=True,
+                 monitor='val_loss', mode='min', save_final_model=False,
+                 ):
+        self.early_stopping = EarlyStopping(
+            patience=patience,
+            delta=delta,
+            checkpoint_path=checkpoint_path,
+            verbose=verbose,
+            restore_best_weights=restore_best_weights,
+            monitor=monitor,
+            mode=mode
+        )
+        self.checkpoint_path = checkpoint_path
+        self.save_final_model = save_final_model
+
+    def on_training_begin(self, state):
+        state['early_stop'] = False
+        state['best_epoch'] = 0
+        state['best_model_path'] = self.checkpoint_path
+
+    def on_validation_end(self, state):
+        epoch = state.get('epoch', None)
+        metric_name = self.early_stopping.monitor
+        # print(state.keys())
+        # print(state.get('val_metrics'))
+
+        if metric_name.startswith('val_metrics'):
+            key = metric_name.split('.')[-1]
+            metric_value = state['val_metrics'].get(key)
+        elif metric_name.startswith('train_metrics'):
+            key = metric_name.split('.')[-1]
+            metric_value = state['train_metrics'].get(key)
+        else:
+            metric_value = state.get(metric_name)
+
+        if metric_value is None:
+            raise ValueError(f"Metric '{metric_name}' not found in state dictionary.")
+
+        stop = self.early_stopping(metric_value, state['model'], epoch=epoch)
+
+        state['early_stop'] = stop
+        state['best_epoch'] = self.early_stopping.best_epoch
+        state['best_val_loss'] = self.early_stopping.best_value
+
+    def on_training_end(self, state):
+        model = state.get('model')
+        if model:
+            self.early_stopping.restore_best_weights(model)
+            print(f"Training finished. Best epoch: {self.early_stopping.best_epoch}, "
+                  f"Best {self.early_stopping.monitor}: {self.early_stopping.best_value:.4f}")
+
+            if self.save_final_model:
+                final_path = self.checkpoint_path.replace(".pt", "_final.pt")
+                torch.save(model.state_dict(), final_path)
+                print(f"Final model saved at {final_path}")

utils/eval_utils.py

@@ -0,0 +1,78 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+from torch import load
+from torch.utils.data import DataLoader
+from typing import Union
+
+from model.epu import EPUCNN
+from utils.omega_parser import EPUCNNParams
+from utils.config_utils import model_cfg_to_epucnn
+
+
+class EPUCNNEval(EPUCNN):
+    def __init__(self, epu_cfg: EPUCNNParams):
+        super().__init__(**model_cfg_to_epucnn(epu_cfg))
+        self.eval()
+
+    # @staticmethod
+    # def initialize_model_from_config(epu_cfg: EPUCNNParams, device: torch.device) -> EPUCNN:
+    #     model = EPUCNN(**model_cfg_to_epucnn(epu_cfg))
+    #     model.to(device)
+    #     return model
+
+    # @staticmethod
+    # def load_ckpt(model: Union[nn.Module, EPUCNN], device: torch.device, ckpt_path: str) -> Union[nn.Module, EPUCNN]:
+    #     state_dict = load(ckpt_path, map_location=device)
+    #     model.load_state_dict(state_dict)
+    #     model.to(device)
+    #     model.eval()
+    #     return model
+
+    def load_ckpt(self, device: torch.device, ckpt_path: str):
+        state_dict = load(ckpt_path, map_location=device)
+        self.load_state_dict(state_dict)
+        self.to(device)
+        self.eval()
+        return self
+
+    # @staticmethod
+    # def get_pretrained_model_from_config(epu_cfg: EPUCNNParams,
+    #                                      device: torch.device,
+    #                                      ckpt_path: str
+    #                                      ):
+    #     model = EPUCNNEval.initialize_model_from_config(epu_cfg, device)
+    #     model = EPUCNNEval.load_ckpt(model, device, ckpt_path)
+    #     return model
+
+
+class InferenceRunnerEPUCNN:
+    def __init__(self, epu_model: Union[EPUCNN, nn.Module], device: torch.device, mode: str = 'binary'):
+        self.epu_model = epu_model
+        self.device = device
+        self.mode = mode
+
+    def predict(self, dataloader: DataLoader, raw_logits=False, return_predictions: bool = False):
+        self.epu_model.eval()
+
+        all_targets = []
+        all_predictions = []
+
+        results = {}
+
+        with torch.no_grad():
+            for batch in dataloader:
+                x, y = batch
+                x = x.to(self.device)
+                y = y.to(self.device, dtype=torch.float32).unsqueeze(1)  # from [bs] to [bs, 1]
+                y_hat = self.epu_model(x, ret_raw_logits=raw_logits)
+
+                all_predictions.append(y_hat.cpu().detach().numpy())
+                if return_predictions:
+                    all_targets.append(y.cpu().detach().numpy())
+
+            results["predictions"] = np.concatenate(all_predictions, axis=0)
+            if return_predictions:
+                results["targets"] = np.concatenate(all_targets, axis=0)
+            return results

utils/metrics.py

@@ -0,0 +1,110 @@
+import numpy as np
+import torch
+import torch.nn as nn
+
+from typing import Optional, Dict, Tuple, Union, Callable
+from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score
+
+
+class EPUMetrics:
+
+    def __init__(self,
+                 mode:              str = "binary",
+                 n_classes:         int = 1,
+                 confidence_level:  float = 0.5,
+                 activation:        Union[nn.Module, Callable] = nn.Sigmoid(),
+                 metrics_config:    Optional[Dict[str, bool]] = None,
+                 ):
+
+        assert mode in ["binary", "multiclass"], "mode should be either 'binary' or 'multiclass'"
+        self.mode = mode
+        
+        self.n_classes = n_classes
+        self.confidence = confidence_level
+
+        self.activation = activation
+
+        default_metrics = {"accuracy": True,  "auc": True}
+        self.metrics_config = metrics_config if metrics_config is not None else default_metrics
+
+        self.avg_method = 'binary' if self.mode == "binary" else 'macro'
+
+    @staticmethod
+    def _to_tensor(x: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
+        if isinstance(x, torch.Tensor):
+            x = x.detach()
+            return x
+        else:
+            return torch.as_tensor(x)
+
+    @staticmethod
+    def _to_numpy(x: torch.Tensor) -> np.array:
+        x = x.detach().cpu().numpy()
+        return x
+
+    def _apply_activation(self, raw_logits: torch.Tensor) -> torch.Tensor:
+        return self.activation(raw_logits)
+
+    def _format_labels(self,
+                       y_true: torch.Tensor,
+                       y_pred: torch.Tensor,
+                       y_prob: torch.Tensor = None) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        if y_prob is None:
+            y_prob = self._apply_activation(y_pred)
+
+        if self.mode == 'binary':
+            y_true = y_true.view(-1)
+            y_prob = y_prob.view(-1)
+            y_pred = (y_prob > self.confidence).to(dtype=torch.int64)
+
+        else:
+            if y_true.ndim > 1:
+                y_true = np.argmax(y_true, axis=1)
+            y_pred = np.argmax(y_pred, axis=1)
+
+        y_true = self._to_numpy(y_true)
+        y_pred = self._to_numpy(y_pred)
+        y_prob = self._to_numpy(y_prob)
+        return y_true, y_pred, y_prob
+
+    @staticmethod
+    def convert_to_onehot(labels: np.ndarray, n_classes: int) -> np.ndarray:
+        if n_classes is None:
+            raise ValueError("Number of classes is not provided")
+        labels = labels.astype(int)
+        labels = np.eye(n_classes)[labels]
+        return labels
+
+    def compute(self,
+                y_true: Union[torch.Tensor, np.ndarray],
+                y_pred: Union[torch.Tensor, np.ndarray],
+                y_prob: Optional[Union[torch.Tensor, np.ndarray]] = None
+                ) -> Dict[str, float]:
+        y_true = self._to_tensor(y_true)
+        y_pred = self._to_tensor(y_pred)
+        y_prob = None if y_prob is None else self._to_tensor(y_prob)
+
+        y_true_np, y_pred_np, y_prob_np = self._format_labels(
+            y_true=y_true,
+            y_pred=y_pred,          # raw_logits
+            y_prob=y_prob)
+
+        metrics = {}
+        if self.metrics_config["accuracy"]:
+            score = accuracy_score(y_true_np, y_pred_np)
+            metrics["accuracy"] = float(np.round(score, 4))
+        if self.metrics_config["auc"]:
+            try:
+                if self.avg_method == "binary":
+                    score = roc_auc_score(y_true_np, y_prob_np)
+                else:
+                    y_true_np = self.convert_to_onehot(y_true_np, n_classes=self.n_classes)
+                    score = roc_auc_score(y_true_np, y_prob_np, multi_class='ovr')
+                metrics["auc"] = score
+            except ValueError:
+                metrics["auc"] = float('nan')       # AUC cannot be calculated
+
+        return metrics
+
+    def update_confidence(self, confidence_level: float) -> None:
+        self.confidence = confidence_level

utils/omega_parser.py

@@ -0,0 +1,109 @@
+from omegaconf import OmegaConf
+from omegaconf import MISSING
+
+from dataclasses import dataclass, field, asdict
+from typing import Any, Dict, Tuple, Optional
+
+
+@dataclass
+class DataLoading:
+    batch_size:         int = 64
+    shuffle:            bool = True
+    num_workers:        int = 0
+    pin_memory:         bool = False
+    persistent_workers: bool = False
+
+
+@dataclass
+class DataPreprocessing:
+    data_parser:    str = MISSING
+    label_mapping:  Dict[str, int] = field(default_factory=dict)
+    resize_dims:    Tuple[int, int] = (128, 128)
+    data_mode:      str = "rgb"                # auto-completed from input_channels in SubnetworkParams
+    medmnist_csv_file:  Optional[str] = None
+
+
+@dataclass
+class DataParams:
+    dataset_path:       str = MISSING
+    images_extension:   str = MISSING
+    data_preprocessing: DataPreprocessing = field(default_factory=DataPreprocessing)
+    data_loading:       DataLoading = field(default_factory=DataLoading)
+
+
+@dataclass
+class TrainingParams:
+    mode:           str = "binary"
+    loss:           str = "binary_cross_entropy"
+    epochs:         int = 10
+    optimizer:      str = "adam"
+    learning_rate:  float = 0.001
+    momentum:       float = 0.0
+    weight_decay:   float = 0.001
+    early_stopping_patience: int = 4
+    early_stopping_monitor: str = "val_loss"
+    early_stopping_mode:    str = "min"
+
+
+@dataclass
+class SubnetworkParams:
+    architecture:       str = "base_one"
+    input_channels:     int = 3
+    # fc_pred_units:    int = 1            # num_classes
+    base_channels:      int = 32
+    fc_hidden_units:    int = 64
+    pred_activation:    str = "tanh"
+
+
+@dataclass
+class EPUCNNParams:
+    num_subnetworks:   int = 4
+    num_classes:       int = 1
+    epu_activation:    str = "sigmoid"
+    subnetwork_config: SubnetworkParams = field(default_factory=SubnetworkParams)
+
+
+@dataclass
+class IUSConfig:
+    model:          EPUCNNParams = field(default_factory=EPUCNNParams)
+    train_params:   TrainingParams = field(default_factory=TrainingParams)
+    data_params:    DataParams = field(default_factory=DataParams)
+    log_dir:            Optional[str] = "./logs"
+    checkpoint_dir:     Optional[str] = "./checkpoints"
+    experiment_name:    Optional[str] = "ius_experiment"
+    timestamp:          Optional[str] = None
+    experiment_saved_folder_name: Optional[str] = None
+
+    @staticmethod
+    def from_yaml(filepath: str) -> "IUSConfig":
+        yaml_cfg = OmegaConf.load(filepath)
+        merged = OmegaConf.merge(OmegaConf.structured(IUSConfig), yaml_cfg)
+        merged = OmegaConf.to_object(merged)
+        if isinstance(merged, IUSConfig):
+            cfg = merged
+        else:               # manually
+            cfg = IUSConfig(**merged)
+        return cfg
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "model": vars(self.model),
+            "train_params": vars(self.train_params),
+            "data_params": vars(self.data_params),
+        }
+
+    def __repr__(self):
+        return OmegaConf.to_yaml(OmegaConf.structured(self), resolve=True)
+
+
+if __name__ == "__main__":
+     cfg_file = "configs/train_config.yaml"
+     cfg = IUSConfig.from_yaml(cfg_file)
+     print(cfg.model)
+     print(cfg.train_params)
+     print(cfg.data_params)
+
+     # yaml_cfg = IUSConfig.from_yaml("checkpoints/test_experiment_unet_0032_20260203_192949/user_config.yaml")
+     # print(yaml_cfg)
+     # print(type(yaml_cfg))
+     # print(yaml_cfg.data_params.test_split)

utils/sanity_utils.py

@@ -0,0 +1,96 @@
+from .omega_parser import IUSConfig
+
+
+class SanityChecker:
+    VALID_MODELS = ["base_one"]
+    VALID_DATA_MODES = ["grayscale", "rgb"]
+    VALID_DATA_PARSERS = ["filename", "folder", "medmnist"]
+
+    def __init__(self, cfg: IUSConfig):
+        self.cfg = cfg
+
+    def model_cfg(self) -> None:
+        architecture = self.cfg.model.subnetwork_config.architecture
+
+        if architecture not in self.VALID_MODELS:
+            raise ValueError(
+                f"EPU-CNN backbone implemented {self.VALID_MODELS}. Original IUS paper was using 'base_one'"
+            )
+
+        num_classes = self.cfg.model.num_classes
+        epu_activation = self.cfg.model.epu_activation
+
+        if epu_activation == "sigmoid" and num_classes > 1:
+            raise ValueError(
+                f"epu_activation = {epu_activation} and num_classes = {num_classes}. "
+                f"For sigmoid epu_activation you have to set num_classes=1"
+            )
+
+        num_subnets = self.cfg.model.num_subnetworks
+        if num_subnets != 4:
+            raise ValueError(
+                "IUS measure was implemented using 4 PFMS (either for grayscale or rgb modalities)"
+                "class PerceptualFeatureMapTransform in data.perceptual_transforms yields 4 PFM representations. "
+            )
+
+    def train_cfg(self) -> None:
+        mode = self.cfg.train_params.mode
+        loss = self.cfg.train_params.loss
+
+        if mode == "binary" and loss != "binary_cross_entropy":
+            raise ValueError(
+                f"In train_params: mode={mode}, loss={loss}. For mode = 'binary' set loss = 'binary_cross_entropy'"
+            )
+
+    def data_preprocessing(self) -> None:
+        data_mode = self.cfg.data_params.data_preprocessing.data_mode
+        channels = self.cfg.model.subnetwork_config.input_channels
+
+        if data_mode not in self.VALID_DATA_MODES:
+            raise ValueError(
+                f"data_mode should be one of {self.VALID_DATA_MODES}. "
+            )
+
+        if channels != 1:
+            raise ValueError(
+                f"In train_params:input_channels={channels} But Perceptual Feature decomposition yields PFMs "
+                f"with 1 output channel. Please set input_channels=1"
+            )
+
+        num_classes = self.cfg.model.num_classes
+        labels = set(self.cfg.data_params.data_preprocessing.label_mapping.values())
+
+        if num_classes > 1 and len(labels) != num_classes:
+            raise ValueError(
+                f"In train_params: num_classes={num_classes}, in label mapping found {labels}. "
+            )
+        if num_classes == 1 and len(labels) != 2:
+            raise ValueError(
+                f"In train_params: num_classes={num_classes}, in label mapping found {labels}. "
+            )
+
+        data_parser = self.cfg.data_params.data_preprocessing.data_parser
+        if data_parser not in self.VALID_DATA_PARSERS:
+            raise ValueError(
+                f"data_parser should be one of {self.VALID_DATA_PARSERS}. "
+                f"Otherwise implement your own data_parser in data.parsers and update utils.sanity_utils"
+            )
+
+    def sanity_check(self) -> None:
+        self.model_cfg()
+        self.train_cfg()
+        self.data_preprocessing()
+        print("All sanity checks passed!")
+
+
+if __name__ == "__main__":
+    from .omega_parser import IUSConfig
+    cfg_file = "configs/train_config.yaml"
+    cfg = IUSConfig.from_yaml(cfg_file)
+
+    # print(cfg.model)
+    # print(cfg.train_params)
+    # print(cfg.data_params)
+
+    checker = SanityChecker(cfg)
+    checker.sanity_check()

utils/tensorboard.py

@@ -0,0 +1,103 @@
+import os
+import time
+import subprocess
+import torch
+import torch.nn as nn
+
+from torch.utils.tensorboard import SummaryWriter
+
+
+class TensorboardLogger(object):
+    def __init__(self, log_dir='runs/experiment'):
+        self.writer = SummaryWriter(log_dir=log_dir)
+
+    def log_scalar(self, tag: str, value: float, step: int):
+        self.writer.add_scalar(tag, value, step)
+
+    # def log_scalars(self, tag: str, values: dict, step: int):
+    #     self.writer.add_scalars(tag, values, step)
+
+    def log_histogram(self, model:nn.Module, step: int):
+        for name, param in model.named_parameters():
+            self.writer.add_histogram(f'weights/{name}', param, step)
+            if param.grad is not None:
+                self.writer.add_histogram(f'gradients/{name}', param.grad, step)
+
+    def log_model_graph(self, model: nn.Module, input_sample: torch.Tensor):
+        self.writer.add_graph(model, input_sample)
+
+    def close(self):
+        self.writer.close()
+
+
+def launch_tensorboard(log_dir: str = "runs", port: int = 6006, open_browser: bool = True):
+    """Automatically launch TensorBoard pointing to log_dir."""
+    if not os.path.exists(log_dir):
+        os.makedirs(log_dir, exist_ok=True)
+    print(f"[TensorBoard] launching at http://localhost:{port}/")
+
+    tb_command = ["tensorboard", f"--logdir={log_dir}", f"--port={port}"]
+    if not open_browser:
+        tb_command.append("--host=127.0.0.1")
+
+    subprocess.Popen(tb_command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+    time.sleep(2)
+
+
+class TensorboardLoggerCallback(object):
+    """
+            Adapter uses TensorboardLogger
+    """
+    def __init__(self,
+                 log_dir: str = 'logs/test_experiment',
+                 log_histograms: bool = False,
+                 launch_tb: bool = True,
+                 tb_port: int = 6006,
+                 open_tb_in_browser: bool = False,
+                 ):
+
+        self.tb_logger = TensorboardLogger(log_dir=log_dir)
+
+        self.log_histograms = log_histograms
+
+        if launch_tb:
+            launch_tensorboard(log_dir=log_dir,
+                               port=tb_port,
+                               open_browser=open_tb_in_browser)
+
+        self.model = None
+
+    def on_train_begin(self, model: nn.Module):
+        self.model = model
+        print("[TensorBoard] Training started, callback init")
+
+    def on_train_end(self):
+        self.tb_logger.close()
+        print("[TensorBoard] Training ended, logs saved")
+
+    def on_epoch_end(self, state):
+        epoch = state.get('epoch', 0)
+
+        # losses
+        self.tb_logger.log_scalar('Loss/train', state.get('train_loss', 0.0), epoch)
+        self.tb_logger.log_scalar('Loss/val', state.get('val_loss', 0.0), epoch)
+
+        # training & validation metrics
+
+        # for metric, value in train_metrics.items():
+        #     self.tb_logger.log_scalar(f'Metrics/train/{metric}', value, epoch)
+        train_metrics = state.get('train_metrics')
+        if train_metrics is not None:
+            # self.tb_logger.log_scalars(f'Metrics/train', train_metrics, epoch)
+            for metric_name, metric_value in train_metrics.items():
+                self.tb_logger.log_scalar(f'Metrics/train/{metric_name}', metric_value, epoch)
+        val_metrics = state.get('val_metrics')
+        if val_metrics is not None:
+            for metric_name, metric_value in val_metrics.items():
+                self.tb_logger.log_scalar(f'Metrics/val/{metric_name}', metric_value, epoch)
+
+        # weights and grads
+        model = state.get('model')
+        if self.log_histograms and model is not None:
+            self.tb_logger.log_histogram(model, epoch)
+

utils/train_utils.py

@@ -0,0 +1,108 @@
+import os
+import json
+
+from typing import Any, List, Union
+from datetime import datetime
+from omegaconf import OmegaConf
+
+from utils.omega_parser import IUSConfig
+
+
+def get_next_experiment_id(log_root: str, base_name: str, pad: int = 4) -> str:
+    """
+    Finds the next zero-padded experiment ID based on existing folders.
+
+    Args:
+        log_root:   Root directory where experiment folders are stored.
+        base_name:  Name of experiment e.g., "experiment_ius"
+        pad:        Number of digits to pad the ID (default: 4).
+
+    Returns:
+        Zero-padded string for next experiment ID, e.g., "0000", "0001".
+    """
+    existing = [
+        d for d in os.listdir(log_root)
+        if os.path.isdir(os.path.join(log_root, d)) and d.startswith(base_name)
+    ]
+
+    # Extract numeric ID
+    existing_ids = []
+    for d in existing:
+        try:
+            num_str = d.split("_")[-3]
+            existing_ids.append(int(num_str))
+        except (ValueError, IndexError):
+            continue
+
+    next_id = max(existing_ids, default=-1) + 1
+    return str(next_id).zfill(pad)
+
+
+def create_experiment_folder(log_root: str, model: str, experiment: str, timestamp: str = None) -> str:
+    """
+    Creates experiment folder with professional naming:
+    experiment_0000_YYYYMMDD_HHMMSS
+
+    Args:
+        log_root: Root directory for logs.
+        task: Downstream task name.
+        model: Model architecture name.
+        dataset: Dataset name.
+
+    Returns:
+        experiment_name
+        @param timestamp:
+    """
+    base_name = f"{experiment}_{model}"
+    exp_id = get_next_experiment_id(log_root, base_name)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") if timestamp is None else timestamp
+    experiment_name = f"{base_name}_{exp_id}_{timestamp}"
+
+    return experiment_name
+
+
+# def create_output_folders(log_root: str, ckpt_root: str, experiment_name: str) -> None:
+#     os.makedirs(os.path.join(log_root, experiment_name), exist_ok=True)
+#     os.makedirs(os.path.join(ckpt_root, experiment_name), exist_ok=True)
+def create_output_folders(folder_list: Union[List[str], str]) -> None:
+    if isinstance(folder_list, str):
+        folder_list = [folder_list]
+    for folder_name in folder_list:
+        os.makedirs(folder_name, exist_ok=True)
+
+
+def save_config_to_output_folder(out_folder: str, cfg: IUSConfig, cfg_filename="user_config.yaml") -> None:
+    config_path = os.path.join(out_folder, cfg_filename)
+    OmegaConf.save(config=OmegaConf.structured(cfg), f=config_path)
+
+
+def update_experiment_metadata(cfg: Any,
+                               experiment_name: str = None,
+                               timestamp: str = None) -> None:
+    """
+    Update experiment metadata fields. Uses given timestamp if provided,
+    otherwise does not override.
+    """
+    if timestamp:
+        cfg.timestamp = timestamp
+
+    if experiment_name and not getattr(cfg, "experiment_name", None):
+        cfg.experiment_name = experiment_name
+
+    if hasattr(cfg, "train_params"):
+        train_params = cfg.train_params
+        train_params.experiment_name = experiment_name
+    if hasattr(cfg, "model"):
+        model_cfg = cfg.model
+        model_cfg.experiment_name = experiment_name
+
+    cfg.experiment_saved_folder_name = experiment_name
+
+
+def save_to_json(diction_to_save: dict, saving_path: str) -> None:
+    try:
+        with open(saving_path, "w") as f:
+            json.dump(diction_to_save, f, indent=4)
+        print(f"Saved json file at {saving_path}")
+    except Exception as e:
+        print(f"Failed to save json file at {saving_path}: {e}")

utils/trainer.py

@@ -0,0 +1,206 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+from tqdm import tqdm
+from typing import List, Any, Optional, Tuple, Dict
+
+# mine
+from data.dataloader import DataLoader
+
+
+class EPUTrainer:
+    def __init__(self,
+                 model:         nn.Module,
+                 device:        torch.device,
+                 optimizer:     optim.Optimizer,
+                 criterion:     nn.Module,
+                 epochs:        int,
+                 train_loader:  DataLoader,
+                 val_loader:    Optional[DataLoader] = None,
+                 callbacks:     Optional[List[object]] = None,
+                 metrics:       Optional = None,
+                 checkpoint_dir:    Optional[str] = None,
+                 ):
+        self.model = model
+        self.val_loader = val_loader
+        self.train_loader = train_loader
+
+        self.device = device
+        self.epochs = epochs
+        self.optimizer = optimizer
+        self.criterion = criterion
+        self.callbacks = callbacks or []
+        self.checkpoint_dir = checkpoint_dir
+
+        self.metrics_fun = metrics
+        # if self.metrics_fun is None:
+
+        # init values
+        self.best_metric = float("inf")
+        self.best_model_path = None
+        self.history = []
+
+        self.state = {"model": self.model,
+                      "epoch": 0,
+                      "early_stop": False,
+                      }
+
+    def train(self):
+        self.model.to(self.device)
+
+        self._on_training_begin()
+
+        for epoch in range(self.epochs):
+            self.state["epoch"] = epoch
+            self._on_epoch_begin()
+
+            train_loss, train_metrics = self._train_one_epoch()
+            val_loss, val_metrics = self._validate_epoch()
+
+            self.history.append({"epoch": epoch,
+                                 "train_loss": train_loss,
+                                 "val_loss": val_loss,
+                                 "train_metrics": train_metrics,
+                                 "val_metrics": val_metrics,}
+                                )
+
+            self._on_epoch_end(train_loss, train_metrics, val_loss, val_metrics)
+            self._on_validation_end()
+
+            if self.state.get("early_stop", False):
+                print("Early stopping triggered.")
+                break
+
+        self._on_training_end()
+        # self._export_metrics_to_json()
+
+    def _train_one_epoch(self) -> Tuple[float, Dict[str, float]]:
+        self.model.train()
+        running_loss = 0.0
+        predictions, ground_truth = [], []
+
+        for i, sample in enumerate(tqdm(self.train_loader, desc=f"Training Epoch {self.state['epoch'] + 1}")):
+            x, y = sample
+            x = x.to(self.device)
+            y = y.to(self.device, dtype=torch.float32).unsqueeze(1)     # from [bs] to [bs, 1]
+
+            self.optimizer.zero_grad()
+
+            y_hat = self.model(x, ret_raw_logits=True)                 # w/o EPU activation -applied internally in loss
+            loss = self.criterion(y_hat, y)
+
+            loss.backward()
+            self.optimizer.step()
+
+            running_loss += loss.item()
+            predictions.append(y_hat.detach().cpu())
+            ground_truth.append(y.detach().cpu())
+
+            for callback in self.callbacks:
+                if hasattr(callback, "on_batch_end"):
+                    callback.on_batch_end(
+                        {**self.state,
+                         "batch": i,
+                         "loss": loss.item()}
+                    )
+
+        avg_loss = running_loss / len(self.train_loader)
+
+        metrics = {}
+        if self.metrics_fun is not None:
+            metrics = self.metrics_fun.compute(
+                y_true=torch.cat(ground_truth, axis=0),
+                y_pred=torch.cat(predictions, axis=0)
+            )
+        return avg_loss, metrics
+
+    def _validate_epoch(self) -> Tuple[float, Dict[str, float]]:
+        if self.val_loader is None:
+            return 0.0, {}
+
+        self.model.eval()
+        total_loss = 0
+        predictions, ground_truths = [], []
+
+        with torch.no_grad():
+            for sample in tqdm(self.val_loader, desc="Validating"):
+                x, y = sample
+                x = x.to(self.device)
+                y = y.to(self.device, dtype=torch.float32).unsqueeze(1)  # from [bs] to [bs, 1]
+                y_hat = self.model(x, ret_raw_logits=True)
+                loss = self.criterion(y_hat, y)
+
+                total_loss += loss.item()
+                predictions.append(y_hat.detach().cpu())
+                ground_truths.append(y.detach().cpu())
+
+        avg_loss = total_loss / len(self.val_loader)
+        metrics = {}
+        if self.metrics_fun is not None:
+            metrics = self.metrics_fun.compute(
+                y_true=torch.cat(ground_truths, axis=0),
+                y_pred=torch.cat(predictions, axis=0)
+            )
+
+        return avg_loss, metrics
+
+    def _on_training_begin(self):
+        for callback in self.callbacks:
+            if hasattr(callback, "on_training_begin"):
+                callback.on_training_begin(self.state)
+
+    def _on_epoch_begin(self):
+        for callback in self.callbacks:
+            if hasattr(callback, "on_epoch_begin"):
+                callback.on_epoch_begin(self.state)
+
+    def _on_epoch_end(self, train_loss, train_metrics, val_loss, val_metrics):
+        # update state
+        self.state.update(
+            {"train_loss": train_loss,
+             "val_loss": val_loss,
+             "train_metrics": train_metrics,
+             "val_metrics": val_metrics,
+             }
+        )
+        # print losses
+        print(f"Epoch {self.state['epoch'] + 1} | "
+              f"Train loss: {train_loss:.4f} | Validation Loss: {val_loss:.4f}")
+
+        # print metrics
+        if train_metrics is not None:
+            train_metrics_str = " | ".join([f"{k}: {v:.4f}" for k, v in train_metrics.items()])
+            print(f"Train metrics:\t\t {train_metrics_str}")
+        if val_metrics:
+            val_metrics_str = " | ".join([f"{k}: {v:.4f}" for k, v in val_metrics.items()])
+            print(f"Validation metrics:\t {val_metrics_str}")
+
+        # exec callbacks
+        for callback in self.callbacks:
+            if hasattr(callback, "on_epoch_end"):
+                callback.on_epoch_end(self.state)
+
+    def _on_validation_end(self,):
+        for callback in self.callbacks:
+            if hasattr(callback, "on_validation_end"):
+                # print(self.state)
+                callback.on_validation_end(self.state)
+
+    def _on_training_end(self):
+        for callback in self.callbacks:
+            if hasattr(callback, "on_training_end"):
+                callback.on_training_end(self.state)
+
+    def get_model(self) -> torch.nn.Module:
+        return self.model
+
+    def get_metrics(self):
+        return self.metrics_fun
+
+    # def _export_metrics_to_json(self):
+    #     if self.checkpoint_dir is not None:
+    #         metrics_path = os.path.join(self.checkpoint_dir, "metrics.json")
+    #         with open(metrics_path, "w") as f:
+    #             json.dump(self.history, f, indent=4)
+    #         print(f"Metrics exported to {metrics_path}")