Keep connected regions of img_pos and img_neg that include the positive points and negative points separately. The function is used for merging automatic results with interactive results in VISTA3D. Args: img_pos: bool type tensor, shape [B, 1, H, W, D], where B means the f
(
img_pos: NdarrayTensor,
img_neg: NdarrayTensor,
point_coords: NdarrayTensor,
point_labels: NdarrayTensor,
pos_val: Sequence[int] = (1, 3),
neg_val: Sequence[int] = (0, 2),
margins: int = 3,
)
| 1198 | |
| 1199 | |
| 1200 | def keep_merge_components_with_points( |
| 1201 | img_pos: NdarrayTensor, |
| 1202 | img_neg: NdarrayTensor, |
| 1203 | point_coords: NdarrayTensor, |
| 1204 | point_labels: NdarrayTensor, |
| 1205 | pos_val: Sequence[int] = (1, 3), |
| 1206 | neg_val: Sequence[int] = (0, 2), |
| 1207 | margins: int = 3, |
| 1208 | ) -> NdarrayTensor: |
| 1209 | """ |
| 1210 | Keep connected regions of img_pos and img_neg that include the positive points and |
| 1211 | negative points separately. The function is used for merging automatic results with interactive |
| 1212 | results in VISTA3D. |
| 1213 | |
| 1214 | Args: |
| 1215 | img_pos: bool type tensor, shape [B, 1, H, W, D], where B means the foreground masks from a single 3D image. |
| 1216 | img_neg: same format as img_pos but corresponds to negative points. |
| 1217 | pos_val: positive point label values. |
| 1218 | neg_val: negative point label values. |
| 1219 | point_coords: the coordinates of each point, shape [B, N, 3], where N means the number of points. |
| 1220 | point_labels: the label of each point, shape [B, N]. |
| 1221 | margins: include points outside of the region but within the margin. |
| 1222 | """ |
| 1223 | |
| 1224 | cucim_skimage, has_cucim = optional_import("cucim.skimage") |
| 1225 | |
| 1226 | use_cp = has_cp and has_cucim and isinstance(img_pos, torch.Tensor) and img_pos.device != torch.device("cpu") |
| 1227 | if use_cp: |
| 1228 | img_pos_ = convert_to_cupy(img_pos.short()) # type: ignore |
| 1229 | img_neg_ = convert_to_cupy(img_neg.short()) # type: ignore |
| 1230 | label = cucim_skimage.measure.label |
| 1231 | lib = cp |
| 1232 | else: |
| 1233 | if not has_measure: |
| 1234 | raise RuntimeError("skimage.measure required.") |
| 1235 | img_pos_, *_ = convert_data_type(img_pos, np.ndarray) |
| 1236 | img_neg_, *_ = convert_data_type(img_neg, np.ndarray) |
| 1237 | # for skimage.measure.label, the input must be bool type |
| 1238 | if img_pos_.dtype != bool or img_neg_.dtype != bool: |
| 1239 | raise ValueError("img_pos and img_neg must be bool type.") |
| 1240 | label = measure.label |
| 1241 | lib = np |
| 1242 | |
| 1243 | features_pos, _ = label(img_pos_, connectivity=3, return_num=True) |
| 1244 | features_neg, _ = label(img_neg_, connectivity=3, return_num=True) |
| 1245 | |
| 1246 | outs = np.zeros_like(img_pos_) |
| 1247 | for bs in range(point_coords.shape[0]): |
| 1248 | for i, p in enumerate(point_coords[bs]): |
| 1249 | if point_labels[bs, i] in pos_val: |
| 1250 | features = features_pos |
| 1251 | elif point_labels[bs, i] in neg_val: |
| 1252 | features = features_neg |
| 1253 | else: |
| 1254 | # if -1 padding point, skip |
| 1255 | continue |
| 1256 | for margin in range(margins): |
| 1257 | if isinstance(p, np.ndarray): |
searching dependent graphs…