Any alternatives for torch with skimage.feature.peak_local_max and scipy.optimize.linear_sum_assignment

username: thunlok
cooked: <p>Hi all,</p>
<p>I’m working on a PyTorch-based pipeline for optimizing many small gaussian beam arrays using camera feedback. Right now, I have a function that takes a single 2D image (<code>std_int</code>) and:</p>
<ol>
<li>Detects peaks in the image (using <code>skimage.feature.peak_local_max</code>).</li>
<li>Matches the detected peaks of the gaussian beams to a set of target positions via a cost matrix with <code>scipy.optimize.linear_sum_assignment</code>.</li>
<li>Updates weights and phases at the matched positions.</li>
</ol>
<p>I’d like to extend this to support batched processing, where I input a tensor of shape <code>[B, H, W]</code> representing B images in a batch, and process all elements simultaneously on the GPU.</p>
<p>My goals are:</p>
<ol>
<li>
<p>Implement a batched version of peak detection (like <code>peak_local_max</code>) in pure PyTorch so I can stay on the GPU and avoid looping over the batch dimension.</p>
</li>
<li>
<p>Implement a batched version of linear sum assignment to match detected peaks to target points per batch element.</p>
</li>
<li>
<p>Minimize CPU-GPU transfers and avoid Python-side loops over B if possible (though I realize that for Hungarian algorithm, some loop may be unavoidable).</p>
</li>
</ol>
<p>Questions:</p>
<ul>
<li>Are there known implementations of batched peak detection in PyTorch for 2D images?</li>
<li>Is there any library or approach for batched linear assignment (Hungarian or something similar such Jonker-Volgenant) on GPU? Or should I implement an approximation like Sinkhorn if I need differentiability and batching?</li>
<li>How do others handle this kind of batched peak detection + assignment in computer vision or microscopy tasks?</li>
</ul>
<p>Here are my current two functions that I need to update further for batching. I need to remove/update the numpy use in linear_sum_assignment and peak_local_max:</p>
<pre><code class="lang-auto">def match_detected_to_target(detected, target):
    # not sure if needed, but making detected&amp;target torchized
    detected = torch.tensor(detected, dtype=torch.float32)
    target = torch.tensor(target, dtype=torch.float32)

    cost_matrix = torch.cdist(detected, target, p=2)  # Equivalent to np.linalg.norm in numpy
    
    cost_matrix_np = cost_matrix.cpu().numpy()
    
    row_ind, col_ind = linear_sum_assignment(cost_matrix_np)

    return row_ind, col_ind  

def weights(w, target, w_prev, std_int, coordinates_ccd_first, min_distance, num_peaks, phase, device='cpu'):

    target = torch.tensor(target, dtype=torch.float32, device=device)
    std_int = torch.tensor(std_int, dtype=torch.float32, device=device)
    w_prev = torch.tensor(w_prev, dtype=torch.float32, device=device)
    phase = torch.tensor(phase, dtype=torch.float32, device=device)

    coordinates_t = torch.nonzero(target &gt; 0)  
    image_shape = std_int.shape
    ccd_mask = torch.zeros(image_shape, dtype=torch.float32, device=device)  


    for y, x in coordinates_ccd_first:
        ccd_mask[y, x] = std_int[y, x]


    coordinates_ccd = peak_local_max(
        std_int.cpu().numpy(),  
        min_distance=min_distance,
        num_peaks=num_peaks
    )
    coordinates_ccd = torch.tensor(coordinates_ccd, dtype=torch.long, device=device)

    row_ind, col_ind = match_detected_to_target(coordinates_ccd, coordinates_t)

    ccd_coords = coordinates_ccd[row_ind]
    tgt_coords = coordinates_t[col_ind]

    ccd_y, ccd_x = ccd_coords[:, 0], ccd_coords[:, 1]
    tgt_y, tgt_x = tgt_coords[:, 0], tgt_coords[:, 1]

    intensities = std_int[ccd_y, ccd_x]
    ideal_values = target[tgt_y, tgt_x]
    previous_weights = w_prev[tgt_y, tgt_x]

    updated_weights = torch.sqrt(ideal_values/intensities)*previous_weights

    phase_mask = torch.zeros(image_shape, dtype=torch.float32, device=device)
    phase_mask[tgt_y, tgt_x] = phase[tgt_y, tgt_x]

    w[tgt_y, tgt_x] = updated_weights

    return w, phase_mask


    w, masked_phase = weights(w, target_im, w_prev, std_int, coordinates, min_distance, num_peaks, phase, device)

</code></pre>
<p>Any advice and help are greatly appreciated! Thanks!</p>