infer.py

import numpy as np
import torch
from scipy.ndimage import gaussian_filter
from skimage.feature import peak_local_max
from sklearn.cluster import MeanShift, estimate_bandwidth

import numpy as np
from sklearn.cluster import MeanShift
from scipy.ndimage import maximum_filter

def extract_coordinates_meanshift(heatmap, min_confidence=0.25, bandwidth=5):
    # Step 1: Adaptive Local Peak Detection
    adaptive_filter_size = max(3, min(20, int(min(heatmap.shape) * 0.05)))
    max_filtered = maximum_filter(heatmap, size=adaptive_filter_size)
    peaks = (heatmap == max_filtered) & (heatmap >= min_confidence * np.max(heatmap))

    y, x = np.where(peaks)
    if len(x) == 0:
        return np.array([])  # No detections

    points = np.column_stack((x, y))

    # Step 2: Mean-Shift Clustering to separate dense overlapping detections
    ms = MeanShift(bandwidth=bandwidth, bin_seeding=True)
    ms.fit(points)
    cluster_centers = ms.cluster_centers_

    return np.array(cluster_centers)

# Inference function
def infer(model, pil_image, log = None):
    img = np.array(pil_image)  # Convert PIL image to NumPy array
    if img.shape[-1] == 4:
        img = img[:, :, :3]  # Remove alpha channel if present
    height, width, _ = img.shape
    
    transform = transforms.ToTensor()
    img_tensor = transform(img).unsqueeze(0).cuda()
    
    with torch.no_grad():
        heatmap = model(img_tensor).cpu().squeeze().numpy()
    
    points = extract_coordinates_meanshift(heatmap)
    if log:
        cv2.imwrite(os.path.join("log", log + "heatmap" + str(len(points)) + ".png"), heatmap)
        pil_image.save(os.path.join("log", log))
    return points