Imrpove SFD detect and batch_detect

[yassineAlouini]

In progress...

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]])