1adrianb / 1adrianb/face-alignment
Imrpove SFD detect and batch_detect
- Dominant language
- Python
- Stars
- 7.5k
- Forks
- 1.4k
- PR merge metrics
- No merged PRs in 30d
Description
In progress...
```python
import torch
import torch.nn.functional as F
import cv2
import numpy as np
from .bbox import decode # assume decode supports vectorized inputs
def detect(net, img, device):
# Transpose from (H, W, C) to (C, H, W)
img = img.transpose(2, 0, 1)
# Create a batch of 1. Use np.ascontiguousarray to avoid extra copies.
img = np.expand_dims(np.ascontiguousarray(img), 0)
img = torch.from_numpy(img).to(device, dtype=torch.float32)
return batch_detect(net, img, device)
def batch_detect(net, img_batch, device):
"""
Inputs:
- img_batch: a torch.Tensor of shape (Batch size, Channels, Height, Width)
"""
# It is better to set cudnn.benchmark globally (outside the function)
# rather than on every call (if using CUDA).
if 'cuda' in device:
torch.backends.cudnn.benchmark = True
# Make sure img_batch is on the correct device and in float32.
img_batch = img_batch.to(device, dtype=torch.float32)
# Convert RGB (assumed input) to BGR by flipping the channel dimension.
# (Could also use explicit channel indexing like img_batch = img_batch[:, [2,1,0],:,:])
img_batch = img_batch.flip(-3)
# Subtract the mean
mean = torch.tensor([104.0, 117.0, 123.0], device=device).view(1, 3, 1, 1)
img_batch = img_batch - mean
with torch.no_grad():
olist = net(img_batch)
# Apply softmax on all classification outputs. Assuming that every even-index output
# is a classification output:
olist = [F.softmax(o, dim=1) if idx % 2 == 0 else o for idx, o in enumerate(olist)]
# Transfer outputs to the CPU and convert to numpy.
olist = [o.cpu().numpy() for o in olist]
bboxlists = get_predictions(olist, img_batch.size(0))
return bboxlists
def get_predictions(olist, batch_size):
"""
Vectorized version that obtains candidate detections from the network outputs.
It groups detections per batch sample.
Returns a list of arrays, one per image in the batch, where each array is
of shape (N, 5) representing the 4 bounding box coordinates and the final score.
"""
# Create a list to hold detections for every image
detections_by_image = [[] for _ in range(batch_size)]
# Variances used in decoding
variances = [0.1, 0.2]
num_scales = len(olist) // 2
for i in range(num_scales):
# Get classification and regression results for this scale.
ocls = olist[i * 2] # shape: (batch, num_classes, H, W)
oreg = olist[i * 2 + 1] # shape: (batch, 4, H, W)
# Define the stride (note that 2**(i+2) gives 4,8,16,32,...)
stride = 2 ** (i + 2)
# Use vectorized thresholding: obtain all positions (across the batch) with score > 0.05
# Note: np.where returns a tuple (batch_inds, h_inds, w_inds)
batch_inds, h_inds, w_inds = np.where(ocls[:, 1, :, :] > 0.05)
if batch_inds.size == 0:
continue
# Compute the center coordinates based on stride.
axc = stride / 2 + w_inds * stride
ayc = stride / 2 + h_inds * stride
# Each candidate uses the same prior box dimensions at this scale.
priors = np.vstack((
axc,
ayc,
np.full_like(axc, stride * 4),
np.full_like(ayc, stride * 4)
)).T # shape: (N, 4)
# Gather the scores (expand dims for concatenation later)
scores = ocls[batch_inds, 1, h_inds, w_inds][:, None] # shape: (N, 1)
# Gather regression outputs for the same positions.
# Here, indexing is done on every detection: from oreg (batch, 4, H, W)
locs = oreg[batch_inds, :, h_inds, w_inds] # shape: (N, 4)
# Decode the location predictions using the priors and provided variances.
# (Assuming that decode is implemented to work with vectorized inputs.)
boxes = decode(locs, priors, variances) # expected shape: (N, 4)
# Concatenate the boxes with their scores.
detections = np.concatenate((boxes, scores), axis=1) # shape: (N, 5)
# Group detections by the image index
for b, det in zip(batch_inds, detections):
detections_by_image[b].append(det)
# For every image in the batch, convert list of detections into a numpy array.
for i in range(batch_size):
if detections_by_image[i]:
detections_by_image[i] = np.stack(detections_by_image[i], axis=0)
else:
# If no candidates, return an empty array with shape (0, 5)
detections_by_image[i] = np.empty((0, 5))
return detections_by_image
def flip_detect(net, img, device):
# Flips the image horizontally.
img = cv2.flip(img, 1)
b = detect(net, img, device)
bboxlist = np.zeros(b[0].shape) if b[0].size > 0 else np.empty((0, 5))
if bboxlist.size > 0:
# Adjust the bounding boxes to the original (flipped) image coordinates.
bboxlist[:, 0] = img.shape[1] - b[0][:, 2] # x_min
bboxlist[:, 1] = b[0][:, 1] # y_min remains the same
bboxlist[:, 2] = img.shape[1] - b[0][:, 0] # x_max
bboxlist[:, 3] = b[0][:, 3] # y_max remains the same
bboxlist[:, 4] = b[0][:, 4] # score
return bboxlist
def pts_to_bb(pts):
# Converts a set of points to a bounding box
min_xy = np.min(pts, axis=0)
max_xy = np.max(pts, axis=0)
return np.array([min_xy[0], min_xy[1], max_xy[0], max_xy[1]])
```
Contributor guide
No contributing guide indexed for this repository
Research direction
The issue is about improving the SFD detect and batch_detect functions. The provided code shows the current implementation. Look at the face-alignment library's structure, find where these functions are located (likely in a detection module), and understand the existing test suite for detection. The goal is to optimize and vectorize the detection pipeline, possibly improving performance or clarity. Check for any related issues or PRs about SFD detection.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- computer-vision
- Issue type
- Refactor
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100