tomotwin-cryoet/tomotwin/modules/inference/locator.py at ae4bbfc22d307b139078de09c1aad58663f8b38b · MPI-Dortmund/tomotwin-cryoet · 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
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
"""
Copyright (c) 2022 MPI-Dortmund
SPDX-License-Identifier: MPL-2.0

This file is subject to the terms of the Mozilla Public License, Version 2.0 (MPL-2.0).
The full text of the MPL-2.0 can be found at http://mozilla.org/MPL/2.0/.

For files that are Incompatible With Secondary Licenses, as defined under the MPL-2.0,
additional notices are required. Refer to the MPL-2.0 license for more details on your
obligations and rights under this license and for instructions on how secondary licenses
may affect the distribution and modification of this software.
"""

from abc import ABC, abstractmethod
from typing import List

import numpy as np
import pandas as pd
import tqdm


class Locator(ABC):

    @abstractmethod
    def locate(self, map_output : pd.DataFrame) -> List[pd.DataFrame]:
        """
        Uses the tomotwin_map output to locate particles for each target reference
        """

    @staticmethod
    def extract_subclass_df(map: pd.DataFrame) -> List[pd.DataFrame]:
        '''
        Extract a dataframe for each reference
        :param map: Resullt from running map
        '''
        sub_dfs = []
        for i in range(len(map.attrs["references"])):
            sub = map[["X", "Y", "Z", f"d_class_{i}"]]
            sub.attrs["ref_name"] = map.attrs["references"][i]
            sub.attrs["ref_index"] = i
            sub_dfs.append(sub)
        return sub_dfs

    @staticmethod
    def nms(boxes: pd.DataFrame, size: int, nms_threshold=0.6) -> pd.DataFrame:
        # np.array(len(coords))
        distance_column = f"metric_best"
        boxes.sort_values(distance_column, inplace=True)
        boxes.reset_index(drop=True, inplace=True)
        boxes["width"] = size
        boxes["height"] = size
        boxes["depth"] = size
        boxes_data = np.empty(shape=(len(boxes), 7))
        kd_data = boxes[["X", "Y", "Z"]].to_numpy()
        from sklearn.neighbors import KDTree
        tree = KDTree(kd_data)

        for i in range(len(boxes)):
            row = boxes.loc[i]
            boxes_data[i, 0] = row["X"]
            boxes_data[i, 1] = row["Y"]
            boxes_data[i, 2] = row["Z"]
            boxes_data[i, 3] = row["width"]
            boxes_data[i, 4] = row["height"]
            boxes_data[i, 5] = row["depth"]
            boxes_data[i, 6] = 1  # merged

        for i in tqdm.tqdm(range(len(boxes_data)), desc="Non-maximum-supression"):
            box_i = boxes_data[i, :]

            close_indicis = tree.query_radius(kd_data[i:(i + 1), :], r=size)[0]
            if box_i[6] == 0:
                continue
            else:
                ones = np.ones((len(close_indicis), 7))
                boxes_i_rep = ones * box_i

                ious = Locator._bbox_iou_vec_3d(boxes_i_rep, boxes_data[close_indicis])
                iou_mask = np.zeros(len(boxes_data), dtype=bool)
                iou_mask[close_indicis] = ious > nms_threshold
                iou_mask[i] = False

                boxes_data[iou_mask, 6] = 0

        boxes = boxes[boxes_data[:, 6] == 1]
        boxes.reset_index(drop=True, inplace=True)
        return boxes

    @staticmethod
    def _bbox_iou_vec_3d(boxesA: np.array, boxesB: np.array) -> np.array:
        # 0 x
        # 1 y
        # 2 z
        # 3 w
        # 4 h
        # 5 depth

        x1_min = boxesA[:, 0] - boxesA[:, 3] / 2
        x1_max = boxesA[:, 0] + boxesA[:, 3] / 2
        y1_min = boxesA[:, 1] - boxesA[:, 4] / 2
        y1_max = boxesA[:, 1] + boxesA[:, 4] / 2
        z1_min = boxesA[:, 2] - boxesA[:, 5] / 2
        z1_max = boxesA[:, 2] + boxesA[:, 5] / 2

        x2_min = boxesB[:, 0] - boxesB[:, 3] / 2
        x2_max = boxesB[:, 0] + boxesB[:, 3] / 2
        y2_min = boxesB[:, 1] - boxesB[:, 4] / 2
        y2_max = boxesB[:, 1] + boxesB[:, 4] / 2
        z2_min = boxesB[:, 2] - boxesB[:, 5] / 2
        z2_max = boxesB[:, 2] + boxesB[:, 5] / 2

        intersect_w = Locator._interval_overlap_vec(x1_min, x1_max, x2_min, x2_max)
        intersect_h = Locator._interval_overlap_vec(y1_min, y1_max, y2_min, y2_max)
        intersect_depth = Locator._interval_overlap_vec(z1_min, z1_max, z2_min, z2_max)
        intersect = intersect_w * intersect_h * intersect_depth
        union = boxesA[:, 3] * boxesA[:, 4] * boxesA[:, 5] + boxesB[:, 3] * boxesB[:, 4] * boxesB[:,
                                                                                           5] - intersect
        return intersect / union

    @staticmethod
    def _interval_overlap_vec(x1_min, x1_max, x2_min, x2_max):
        intersect = np.zeros(shape=(len(x1_min)))
        cond_a = x2_min < x1_min
        cond_b = cond_a & (x2_max >= x1_min)
        intersect[cond_b] = np.minimum(x1_max[cond_b], x2_max[cond_b]) - x1_min[cond_b]
        cond_c = ~cond_a & (x1_max >= x2_min)
        intersect[cond_c] = np.minimum(x1_max[cond_c], x2_max[cond_c]) - x2_min[cond_c]

        return intersect