API Reference

Complete API documentation for the draw_heatmap package.

accvlab.draw_heatmap.draw_heatmap(heatmaps: torch.Tensor, centers: torch.Tensor, radii: torch.Tensor, heatmap_idxes: torch.Tensor, diameter_to_sigma_factor: float = 6.0, k_scale: float = 1.0) → None

This function expects input data to be on the GPU.

A function that draws Gaussian heatmaps based on centers and radii.

Parameters:

  • heatmaps – The heatmaps to be drawn. This is a tensor of type float32 with shape (num_heatmaps, height, width).

  • centers – The centers of each Gaussian kernel. This is a tensor of type int32 with shape (num_targets, 2).

  • radii – The radii of each Gaussian kernel. This is a tensor of type int32 with shape (num_targets,).

  • heatmap_idxes – The heatmap indices of each Gaussian kernel. This is a tensor of type int32 with shape (num_targets,).

  • diameter_to_sigma_factor – The factor used to convert the diameter to the standard deviation of the Gaussian kernel. The default value is 6.

  • k_scale – The scale applied to the entries within the Gaussian kernel. The default value is 1.

accvlab.draw_heatmap.draw_heatmap_batched(heatmap, centers, radii, diameter_to_sigma_factor=6.0, k_scale=1.0, labels=None)[source]

This function expects input data to be on the GPU.

Draws heatmaps for a batch of samples.

Parameters:

  • heatmap (Tensor) – Tensor of shape (batch_size, height, width) when labels is None. Otherwise with shape (batch_size, max_num_classes, height, width). The heatmap will be modified in place.

  • centers (RaggedBatch) – RaggedBatch of shape (batch_size, max_num_targets, 2). The centers of the heatmaps to draw. max_num_targets is the maximum number of targets across the batch.

  • radii (RaggedBatch) – RaggedBatch of shape (batch_size, max_num_targets). The radii of the heatmaps to draw. max_num_targets is the maximum number of targets across the batch.

  • diameter_to_sigma_factor (float, default: 6.0) – Factor for converting diameter to sigma.

  • k_scale (float, default: 1.0) – Scale factor for the Gaussian kernel

  • labels (Optional[RaggedBatch], default: None) – RaggedBatch of shape (batch_size, max_num_targets). The labels are denoted as the class index. max_num_targets is the maximum number of targets across the batch. If None, all classes of the sample will be drawn in one heatmap.