STAR-Bench/fabersw-unique.py at main · necst/STAR-Bench · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# /******************************************
#  * MIT License
#  * Copyright (c) 2022 Eleonora D'Arnese, Davide Conficconi, Emanuele Del
#  * Sozzo, Luigi Fusco, Donatella Sciuto, Marco Domenico Santambrogio
#  ******************************************/
"""
Standalone ETNA software/FPGA registration driver.

Historical entry point kept for backwards compatibility: it glob-scans a
folder for ``*b.png`` / ``*a.png`` image pairs and runs the pyramidal ETNA
optimizer on each pair. Prefer ``starbench-main.py`` for new experiments.
"""
import argparse
import glob
import logging
import os
import warnings

from ETNA import (
    EtnaMultiMetric,
    EtnaMultiPowell,
    EtnaMultiOnePlusOne,
)

warnings.filterwarnings("ignore", category=RuntimeWarning, module="numpy")

logger = logging.getLogger(__name__)


def compute_wrapper(args, num_threads: int = 1, image_dimension: int = 512):
    """Iterate over the patients and run the chosen optimizer on each pair."""
    if args.optimizer == 'powell':
        optimizer = EtnaMultiPowell()
    else:
        optimizer = EtnaMultiOnePlusOne()

    for k in range(args.offset, args.patient):
        curr_res = os.path.join("", args.res_path)
        os.makedirs(curr_res, exist_ok=True)

        # Historical naming convention: reference images end in "b.png",
        # moving images end in "a.png".
        CT = sorted(glob.glob(os.path.join(args.ct_path, '*b.png')))
        PET = sorted(glob.glob(os.path.join(args.pet_path, '*a.png')))

        logger.info(f"Reference images found: {CT}")
        logger.info(f"Floating images found:  {PET}")
        assert len(CT) == len(PET), "Reference/floating image counts differ"

        images_per_thread = len(CT) // num_threads
        for i in range(num_threads):
            start = images_per_thread * i
            end = images_per_thread * (i + 1) if i < num_threads - 1 else len(CT)

            metric_component = EtnaMultiMetric(
                ref_size=image_dimension,
                metric=args.metric,
                transform=args.transform,
                exponential=args.exponential,
                interpolation="nearest",
            )
            name = "t%02d" % i
            optimizer.compute(
                CT[start:end], PET[start:end], name, curr_res, i, k,
                metric_component, image_dimension,
            )


def main():
    parser = argparse.ArgumentParser(description='Standalone ETNA SW driver (legacy entry point)')
    parser.add_argument("-pt", "--patient", nargs='?', default=1, type=int,
                        help='Number of the patient to analyze')
    parser.add_argument("-o", "--offset", nargs='?', default=0, type=int,
                        help='Index of the first patient to analyze')
    parser.add_argument("-cp", "--ct_path", nargs='?', default='./',
                        help='Reference images folder')
    parser.add_argument("-pp", "--pet_path", nargs='?', default='./',
                        help='Moving images folder')
    parser.add_argument("-rp", "--res_path", nargs='?', default='./',
                        help='Results folder')
    parser.add_argument("-t", "--thread_number", nargs='?', default=1, type=int,
                        help='Number of parallel workers')
    parser.add_argument("-im", "--image_dimension", nargs='?', default=512, type=int,
                        help='Target image dimension (square)')
    parser.add_argument("-mtr", "--metric", nargs='?', default='mi',
                        help="Metric: one of {mi, mse, cc, prz}")
    parser.add_argument("-tx", "--transform", nargs='?', default='',
                        help='Transform family (legacy, unused)')
    parser.add_argument("-exp", "--exponential", action='store_true',
                        help='Use the exponential variant of the metric (MI / PRZ)')
    parser.add_argument("-opt", "--optimizer", nargs='?', default='powell',
                        help='Optimizer: powell or oneplusone')

    args = parser.parse_args()
    logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s')
    logger.info(args)

    compute_wrapper(args, args.thread_number, args.image_dimension)
    logger.info(f"ETNA SW-{args.optimizer} run complete.")


if __name__ == "__main__":
    main()