convert_seg2det.py 10 KB
Newer Older
mibaumgartner's avatar
mibaumgartner committed
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
"""
Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

   http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import argparse
import sys
from datetime import datetime
from itertools import repeat
from multiprocessing import Pool
from pathlib import Path
from typing import Sequence

import numpy as np
import SimpleITK as sitk
27
from hydra import initialize_config_module
mibaumgartner's avatar
mibaumgartner committed
28
29
30
from loguru import logger
from tqdm import tqdm

mibaumgartner's avatar
core  
mibaumgartner committed
31
from nndet.core.boxes import box_size_np
32
from nndet.io import save_json
mibaumgartner's avatar
mibaumgartner committed
33
from nndet.io.transforms.instances import get_bbox_np
34
from nndet.io.itk import load_sitk, load_sitk_as_array
mibaumgartner's avatar
mibaumgartner committed
35
from nndet.utils.config import compose
mibaumgartner's avatar
mibaumgartner committed
36
from nndet.utils.check import env_guard
37
from nndet.utils.clustering import seg_to_instances
mibaumgartner's avatar
mibaumgartner committed
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59


def prepare_detection_label(case_id: str,
                            label_dir: Path,
                            things_classes: Sequence[int],
                            stuff_classes: Sequence[int],
                            min_size: float = 0,
                            min_vol: float = 0,
                            ):
    if (label_dir / f"{case_id}.json").is_file():
            logger.info(f"Found existing case {case_id} -> skipping")
            return
    logger.info(f"Processing {case_id}")
    seg_itk = load_sitk(label_dir / f"{case_id}.nii.gz")
    spacing = np.asarray(seg_itk.GetSpacing())[::-1]
    seg = sitk.GetArrayFromImage(seg_itk)

    # prepare stuff information
    stuff_seg = np.zeros_like(seg)
    if stuff_classes:
        for new_class, old_class in enumerate(stuff_classes, start=1):
            stuff_seg[seg == old_class] = new_class
mibaumgartner's avatar
mibaumgartner committed
60
61
62
63
64
        stuff_seg_itk = sitk.GetImageFromArray(stuff_seg)
        stuff_seg_itk.SetOrigin(seg_itk.GetOrigin())
        stuff_seg_itk.SetDirection(seg_itk.GetDirection())
        stuff_seg_itk.SetSpacing(seg_itk.GetSpacing())

mibaumgartner's avatar
mibaumgartner committed
65
66
67
68
69
70
        sitk.WriteImage(stuff_seg_itk, str(label_dir / f"{case_id}_stuff.nii.gz"))

    # prepare things information
    things_seg = np.copy(seg)
    things_seg[stuff_seg > 0] = 0  # remove all stuff classes from segmentation

71
    instances_not_filtered, instances_not_filtered_classes = seg_to_instances(things_seg)
mibaumgartner's avatar
mibaumgartner committed
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
    final_mapping = {}
    if instances_not_filtered.max() > 0:
        boxes = get_bbox_np(instances_not_filtered[None])["boxes"]
        box_sizes = box_size_np(boxes)

        instance_ids = np.unique(instances_not_filtered)
        instance_ids = instance_ids[instance_ids > 0]

        assert len(instance_ids) == len(boxes)
        isotopic_axis = list(range(seg.ndim))
        isotopic_axis.pop(np.argmax(spacing))
        instances = np.zeros_like(instances_not_filtered)

        start_id = 1
        for iid, bsize in zip(instance_ids, box_sizes):
            bsize_world = bsize * spacing
            instance_mask = (instances_not_filtered == iid)
            instance_vol = instance_mask.sum()

            if all(bsize_world[isotopic_axis] > min_size) and (instance_vol > min_vol):
                instances[instance_mask] = start_id
93
                semantic_class = instances_not_filtered_classes[int(iid)]
mibaumgartner's avatar
mibaumgartner committed
94
95
96
97
98
                final_mapping[start_id] = things_classes.index(semantic_class)
                start_id += 1
    else:
        instances = np.zeros_like(instances_not_filtered)

mibaumgartner's avatar
mibaumgartner committed
99
100
101
102
103
    final_instances_itk = sitk.GetImageFromArray(instances)
    final_instances_itk.SetOrigin(seg_itk.GetOrigin())
    final_instances_itk.SetDirection(seg_itk.GetDirection())
    final_instances_itk.SetSpacing(seg_itk.GetSpacing())

mibaumgartner's avatar
mibaumgartner committed
104
105
106
107
108
109
    sitk.WriteImage(final_instances_itk, str(label_dir / f"{case_id}.nii.gz"))
    save_json({"instances": final_mapping}, label_dir / f"{case_id}.json")

    sitk.WriteImage(seg_itk, str(label_dir / f"{case_id}_orig.nii.gz"))


mibaumgartner's avatar
mibaumgartner committed
110
111
@env_guard
def main():
mibaumgartner's avatar
mibaumgartner committed
112
113
114
115
116
117
118
119
120
121
122
123
124
125
    """
    This script converts a semantic segmentation dataset into an instance
    segmentation dataset by using connected components on the labels.
    To account for separated pixels inside the annotations only annotations
    with a specified minimal size are converted into objects.
    
    The data needs to be in the same format as in nnunet: images
    stay the same, labels will be semantic segmentations.

    ============================================================================
    ================================IMPORTANT==================================+
    ============================================================================  
    Needs additional information from dataset.json/.yaml:
        `seg2det_stuff`: these are classes which are interpreted semantically
126
127
            (stuff classes are experimental and will probably changed in
             the future)
mibaumgartner's avatar
mibaumgartner committed
128
129
130
        `seg2det_things`: these are classes which are interpreted as instances
         Both entries should be lists with the indices of the respective
         classes where the position will determine its new class
131
         (currently only one classes is supported here)
mibaumgartner's avatar
mibaumgartner committed
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
         e.g.
            `seg2det_stuff`: [2,] -> remap class 2 from semantic segmentation
                to new stuff class 1 (stuff classes start at one)
            `seg2det_things`: [1, 3] -> remap class 1 and 3 from semantic
                segmentation to new things classes 0 and 1, respectively
            `min_size`: minimum size in mm of objects in the isotropic axis (default 0)
            `min_vol`: minimum volume of instances in pixels (default 0)
    ============================================================================

    The segmentation labels will be splitted into things (classes to detect)
    and  stuff classes (additional segmentation labels) and will be saved
    as separate files.
    """
    parser = argparse.ArgumentParser()
    parser.add_argument('tasks', type=str, nargs='+',
                        help="Single or multiple task identifiers to process consecutively",
                        )
    parser.add_argument('--overwrite', action='store_true')
    parser.add_argument('-o', '--overwrites', type=str, nargs='+',
                        help="overwrites for config file",
                        required=False,
                        )
    parser.add_argument('--volume_ranking',
                        help="Create a ranking of instances based on their volume",
                        action='store_true',
                        )
158
    parser.add_argument('--num_processes', type=int, default=4, required=False,
mibaumgartner's avatar
mibaumgartner committed
159
                        help="Number of processes to use for conversion. Default 4.")
mibaumgartner's avatar
mibaumgartner committed
160
161
162
163
164
165

    args = parser.parse_args()
    tasks = args.tasks
    ov = args.overwrites
    overwrite = args.overwrite
    do_volume_ranking = args.volume_ranking
166
    num_processes = args.num_processes
mibaumgartner's avatar
mibaumgartner committed
167
168
169
170
    initialize_config_module(config_module="nndet.conf")

    for task in tasks:
        cfg = compose(task, "config.yaml", overrides=ov if ov is not None else [])
mibaumgartner's avatar
mibaumgartner committed
171
        print(cfg)
mibaumgartner's avatar
mibaumgartner committed
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195

        splitted_dir = Path(cfg["host"]["splitted_4d_output_dir"])

        logger.remove()
        logger.add(sys.stdout, level="INFO")
        logger.add(splitted_dir / "convert_seg2det.log", level="DEBUG")
        logger.info(f"+++++ Running covnersion: {datetime.now()} +++++")
        logger.info(f"Running min_size {cfg['data'].get('min_size', 0)} and "
                    f"min_vol {cfg['data'].get('min_vol', 0)}")

        for postfix in ["Tr", "Ts"]:
            label_dir = splitted_dir / f"labels{postfix}"
            case_ids = [f.name[:-7] for f in label_dir.glob("*.nii.gz")]
            logger.info(f"Found {len(case_ids)} cases for conversion with postfix {postfix}.")

            # for cid in case_ids:
            #     prepare_detection_label(case_id=cid,
            #                             label_dir=label_dir,
            #                             stuff_classes=cfg["data"]["seg2det_stuff"],
            #                             things_classes=cfg["data"]["seg2det_things"],
            #                             min_size=cfg["data"].get("min_size", 0),
            #                             min_vol=cfg["data"].get("min_vol", 0),
            #                             )

196
            with Pool(processes=num_processes) as p:
mibaumgartner's avatar
mibaumgartner committed
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
                p.starmap(prepare_detection_label, zip(
                    case_ids,
                    repeat(label_dir),
                    repeat(cfg["data"]["seg2det_things"]),
                    repeat(cfg["data"]["seg2det_stuff"]),
                    repeat(cfg["data"].get("min_size", 0)),
                    repeat(cfg["data"].get("min_vol", 0)),
                ))

        if do_volume_ranking:
            for postfix in ["Tr", "Ts"]:
                if (label_dir := splitted_dir / f"labels{postfix}").is_dir():
                    ranking = []
                    for case_id in tqdm([f.stem for f in label_dir.glob("*.json")]):
                        instances = load_sitk_as_array(label_dir / f"{case_id}.nii.gz")[0]
                        instance_ids, instance_counts = np.unique(instances, return_counts=True)
                        cps = [np.argwhere(instances == iid)[0].tolist() for iid in instance_ids[1:]]
                        assert len(instance_ids) - 1 == len(cps)
                        tmp = [{"case_id": str(case_id), "instance_id": int(iid),
                                "vol": int(vol), "cp": list(cp)[::-1]}
                               for iid, vol, cp in zip(instance_ids[1:], instance_counts[1:], cps)]
                        ranking.extend(tmp)
                    ranking = sorted(ranking, key=lambda x: x["vol"])
                    save_json(ranking, splitted_dir / f"volume_ranking_{postfix}.json")
                else:
                    logger.info(f"Did not find dir {label_dir} for volume ranking")
mibaumgartner's avatar
mibaumgartner committed
223
224
225
226


if __name__ == '__main__':
    main()