From 17ae697fc063e5123da432ae3bf519a84ddbf102 Mon Sep 17 00:00:00 2001 From: Benjamin Date: Thu, 15 May 2025 19:40:47 +0200 Subject: [PATCH 1/2] fix transport linking --- 1.0.0 | 5 ++ arcos4py/tools/_detect_events.py | 80 +++++++----------- .../pix/4_colliding_transportation.tif | Bin 1806 -> 4122 bytes 3 files changed, 37 insertions(+), 48 deletions(-) create mode 100644 1.0.0 diff --git a/1.0.0 b/1.0.0 new file mode 100644 index 0000000..2d8d6fd --- /dev/null +++ b/1.0.0 @@ -0,0 +1,5 @@ +Collecting poetry-core + Using cached poetry_core-2.1.3-py3-none-any.whl.metadata (3.5 kB) +Using cached poetry_core-2.1.3-py3-none-any.whl (332 kB) +Installing collected packages: poetry-core +Successfully installed poetry-core-2.1.3 diff --git a/arcos4py/tools/_detect_events.py b/arcos4py/tools/_detect_events.py index 7ff6e32..049968a 100644 --- a/arcos4py/tools/_detect_events.py +++ b/arcos4py/tools/_detect_events.py @@ -197,12 +197,13 @@ def brute_force_linking( return cluster_labels, max_cluster_label + @njit(parallel=True) def _compute_filtered_distances(current_coords, memory_coords): n, m = len(current_coords), len(memory_coords) distances = np.empty((n, m)) for i in prange(n): - for j in prange(m): + for j in range(m): distances[i, j] = np.sum((current_coords[i] - memory_coords[j]) ** 2) return np.sqrt(distances) @@ -229,70 +230,53 @@ def transportation_linking( cost_threshold: float = 0, **kwargs: Dict[str, Any], ) -> Tuple[np.ndarray, int]: - """Optimized transportation linking of clusters across frames, using a pre-constructed sklearn KDTree. - - Args: - cluster_labels (np.ndarray): The cluster labels for the current frame. - cluster_coordinates (np.ndarray): The cluster coordinates for the current frame. - memory_cluster_labels (np.ndarray): The cluster labels for previous frames. - memory_coordinates (np.ndarray): The coordinates for previous frames. - memory_kdtree (KDTree): Pre-constructed sklearn KDTree for memory coordinates. - epsPrev (float): Frame-to-frame distance, used to connect clusters across frames. - max_cluster_label (int): The maximum label for clusters. - reg (float): Entropy regularization parameter for Sinkhorn algorithm. - reg_m (float): Marginal relaxation parameter for unbalanced OT. - cost_threshold (float): Threshold for filtering low-probability matches. - **kwargs: Additional keyword arguments. + """ + Optimized pixel-wise transportation linking of clusters across frames. - Returns: - Tuple[np.ndarray, int]: Updated cluster labels and the maximum cluster label. + Uses unbalanced OT to assign each current pixel to a previous pixel within epsPrev. """ - # Find neighbors within the maximum allowed distance (epsPrev) - indices = memory_kdtree.query_radius(cluster_coordinates, r=epsPrev) + # Find all memory indices within epsPrev of any current pixel + neighbors = memory_kdtree.query_radius(cluster_coordinates, r=epsPrev) - if all(len(ind) == 0 for ind in indices): + if all(len(ind) == 0 for ind in neighbors): max_cluster_label += 1 - return np.full_like(cluster_labels, max_cluster_label), max_cluster_label + return np.full(cluster_labels.shape, max_cluster_label, dtype=int), max_cluster_label - # Prepare indices of valid points - valid_mask = np.array([len(ind) > 0 for ind in indices]) - current_indices = np.arange(len(indices))[valid_mask] - memory_indices = np.array([ind[0] for ind in indices if len(ind) > 0]) - - if len(current_indices) == 0: + valid_mem_idx = np.unique(np.concatenate([ind for ind in neighbors if len(ind) > 0])) + if valid_mem_idx.size == 0: max_cluster_label += 1 - return np.full_like(cluster_labels, max_cluster_label), max_cluster_label - - # Compute distance matrix for valid pairs - filtered_distances = _compute_filtered_distances( - cluster_coordinates[current_indices], memory_coordinates[memory_indices] - ) + return np.full(cluster_labels.shape, max_cluster_label, dtype=int), max_cluster_label - # Uniform distribution on the valid points - a = np.ones(len(current_indices)) / len(current_indices) - b = np.ones(len(memory_indices)) / len(memory_indices) + # Build cost matrix between each current pixel and each candidate memory pixel + curr_coords = cluster_coordinates + mem_coords = memory_coordinates[valid_mem_idx] + cost_matrix = _compute_filtered_distances(curr_coords, mem_coords) - # Solve the unbalanced OT problem - ot_plan = ot.unbalanced.sinkhorn_unbalanced(a, b, filtered_distances, reg, reg_m) + # Uniform distributions + n_curr = curr_coords.shape[0] + n_mem = mem_coords.shape[0] + a = np.ones(n_curr) / n_curr + b = np.ones(n_mem) / n_mem - # Propagate cluster id from previous frame - matches = np.argmax(ot_plan, axis=1) + # Solve unbalanced OT + ot_plan = ot.unbalanced.sinkhorn_unbalanced(a, b, cost_matrix, reg, reg_m) - # Set matches to -1 if the cost is too high - matches[ot_plan[np.arange(len(matches)), matches] < cost_threshold] = -1 + # Determine best assignment for each current pixel + best_mem = np.argmax(ot_plan, axis=1) + probs = ot_plan[np.arange(n_curr), best_mem] + best_mem[probs < cost_threshold] = -1 - new_cluster_labels = _assign_labels( - matches, current_indices, memory_indices, memory_cluster_labels, cluster_labels.size - ) + new_cluster_labels = np.full(n_curr, -1, dtype=int) + for i, m in enumerate(best_mem): + if m != -1: + new_cluster_labels[i] = int(memory_cluster_labels[valid_mem_idx[m]]) - # Assign new labels to unmatched clusters if np.any(new_cluster_labels == -1): max_cluster_label += 1 new_cluster_labels[new_cluster_labels == -1] = max_cluster_label return new_cluster_labels, max_cluster_label - @dataclass class Memory: """Memory class for retaining coordinates and cluster IDs over a specified number of time points. @@ -1227,7 +1211,7 @@ def link(self, input_coordinates: np.ndarray) -> None: linked_cluster_ids = self._update_id(original_cluster_ids, coordinates) if self._remove_small_clusters: - final_cluster_ids = self._apply_remove_small_clusters(linked_cluster_ids, original_cluster_ids) + linked_cluster_ids = self._apply_remove_small_clusters(linked_cluster_ids, original_cluster_ids) # Apply stable merges and splits final_cluster_ids = self._apply_stable_merges_splits(linked_cluster_ids, original_cluster_ids) diff --git a/tests/testdata/pix/4_colliding_transportation.tif b/tests/testdata/pix/4_colliding_transportation.tif index 9f47a228dd35794fc2db4c64073563b75754f9f2..abf7511cf656b439ae4f2dbe345ad048b987073a 100644 GIT binary patch literal 4122 zcmeHKy-EW?5T3n@QT&NoXh55gLV`#j2th2YO`BFMED{V@1WY4h;}ck0`wTvYPmt0w zwY4*MH{9IquHJ=R#o!#v&V4gGH~a0}G1u!I!U_P?08|AOY8yC|&;u1l_!i<}qLW`$ z=rt8U9kEv&+sRs5v?F|WMa+7l_t9rhEk*ICxOP9$-%)RgwH=JVj@S&x(;C2i^Y;2` zJZPRaE>4abonso84`|E?WhuvGG#m~`H-q;0F5~xqjHu7WW+L4x%PFCznE5CbF74kTBS z8VWF(vrj}% z8e5kkqoP_uTEmIN$p=LiPvhWb=15{+@rXe*h(Qw+j#TjD0w%(-OROI+0nNhYN09#% zcqAkp1f&n>Hv5|D8LF|VnP{f5rETD_G|F(BcqBEE!7+znJI}5x28CS=wp8+812bX& zSrZwDdw^!)^WOp%1EU57wT8}v88f)lB$z#nl{Zd2>UfC5vM5DDVnV6nnWl&nOk2g> zoOdR4wKB43GBe0LV=$(YPYYNG`&5repUwlCh0CXqR11pF1-uC<430hw!p#>AM12g{ zCZxF~85>McZobf(C3Yl1j7?yPCBAbJS@?Vi_Md3NkwgP! zV5VbITOnZFCSqyC!=vNSktuehK#a$Fg*wnVK@7@Nia!H3!oE`?(sw~Xv+(&2Qr<|| zG)!b<*~qB2K*89IlTD3HOvWf;>Sad(2JSx$Lez;f26n=}LasgVm7$=Nhs#%>_(GKb E0IuDD(*OVf From ec22f0ec92f37afdd403f567e9c79fdc45049087 Mon Sep 17 00:00:00 2001 From: Benjamin Date: Thu, 15 May 2025 19:55:15 +0200 Subject: [PATCH 2/2] fix lint --- arcos4py/tools/_detect_events.py | 23 +++++++++++++++++------ 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/arcos4py/tools/_detect_events.py b/arcos4py/tools/_detect_events.py index 049968a..d8f5f9d 100644 --- a/arcos4py/tools/_detect_events.py +++ b/arcos4py/tools/_detect_events.py @@ -197,7 +197,6 @@ def brute_force_linking( return cluster_labels, max_cluster_label - @njit(parallel=True) def _compute_filtered_distances(current_coords, memory_coords): n, m = len(current_coords), len(memory_coords) @@ -230,8 +229,19 @@ def transportation_linking( cost_threshold: float = 0, **kwargs: Dict[str, Any], ) -> Tuple[np.ndarray, int]: - """ - Optimized pixel-wise transportation linking of clusters across frames. + """Optimized pixel-wise transportation linking of clusters across frames. + + Arguments: + cluster_labels (np.ndarray): The cluster labels for the current frame. + cluster_coordinates (np.ndarray): The cluster coordinates for the current frame. + memory_cluster_labels (np.ndarray): The cluster labels for previous frames. + memory_coordinates (np.ndarray): The cluster coordinates for previous frames. + memory_kdtree (KDTree): KDTree for the previous frame's clusters. + epsPrev (float): Frame-to-frame distance, used to connect clusters across frames. + max_cluster_label (int): The maximum label for clusters. + reg (float): Entropy regularization parameter for unbalanced OT algorithm. + reg_m (float): Marginal relaxation parameter for unbalanced OT. + cost_threshold (float): Cost threshold for assigning clusters. Uses unbalanced OT to assign each current pixel to a previous pixel within epsPrev. """ @@ -249,14 +259,14 @@ def transportation_linking( # Build cost matrix between each current pixel and each candidate memory pixel curr_coords = cluster_coordinates - mem_coords = memory_coordinates[valid_mem_idx] + mem_coords = memory_coordinates[valid_mem_idx] cost_matrix = _compute_filtered_distances(curr_coords, mem_coords) # Uniform distributions n_curr = curr_coords.shape[0] - n_mem = mem_coords.shape[0] + n_mem = mem_coords.shape[0] a = np.ones(n_curr) / n_curr - b = np.ones(n_mem) / n_mem + b = np.ones(n_mem) / n_mem # Solve unbalanced OT ot_plan = ot.unbalanced.sinkhorn_unbalanced(a, b, cost_matrix, reg, reg_m) @@ -277,6 +287,7 @@ def transportation_linking( return new_cluster_labels, max_cluster_label + @dataclass class Memory: """Memory class for retaining coordinates and cluster IDs over a specified number of time points.