utils.py

import torch
import cv2
import pandas as pd
import numpy as np
from pathlib import Path
import matplotlib.pyplot as plt
from constants import *


def output_tensor_to_boxes(boxes_tensor):
    """
        Converts the YOLO output tensor to list of boxes with probabilites.

        Arguments:
        boxes_tensor -- tensor of shape (S, S, BOX, 5)

        Returns:
        boxes -- list of shape (None, 5)

        Note: "None" is here because you don't know the exact number of selected boxes, as it depends on the threshold. 
        For example, the actual output size of scores would be (10, 5) if there are 10 boxes
    """
    cell_w, cell_h = W/S, H/S
    boxes = []
    
    for i in range(S):
        for j in range(S):
            for b in range(BOX):
                anchor_wh = torch.tensor(ANCHORS[b])
                data = boxes_tensor[i,j,b]
                xy = torch.sigmoid(data[:2])
                wh = torch.exp(data[2:4])*anchor_wh
                obj_prob = torch.sigmoid(data[4])
                
                if obj_prob > OUTPUT_THRESH:
                    x_center, y_center, w, h = xy[0], xy[1], wh[0], wh[1]
                    x, y = x_center+j-w/2, y_center+i-h/2
                    x,y,w,h = x*cell_w, y*cell_h, w*cell_w, h*cell_h
                    box = [x,y,w,h, obj_prob]
                    boxes.append(box)
    return boxes


def plot_img(img, size=(7,7)):
    plt.figure(figsize=size)
    plt.imshow(img)
    plt.show()


def plot_normalized_img(img, std=STD, mean=MEAN, size=(7,7)):
    mean = mean if isinstance(mean, np.ndarray) else np.array(mean)
    std = std if isinstance(std, np.ndarray) else np.array(std)
    plt.figure(figsize=size)
    plt.imshow((255. * (img * std + mean)).astype(np.uint))
    plt.show()
    

def visualize_bbox(img, boxes, thickness=2, color=BOX_COLOR, draw_center=True):
    """
        Draws boxes on the given image.

        Arguments:
        img -- torch.Tensor of shape (3, W, H) or numpy.ndarray of shape (W, H, 3)
        boxes -- list of shape (None, 5)
        thickness -- number specifying the thickness of box border
        color -- RGB tuple of shape (3,) specifying the color of boxes
        draw_center -- boolean specifying whether to draw center or not

        Returns:
        img_copy -- numpy.ndarray of shape(W, H, 3) containing image with bouning boxes
    """
    img_copy = img.cpu().permute(1,2,0).numpy() if isinstance(img, torch.Tensor) else img.copy()
    for box in boxes:
        x,y,w,h = int(box[0]), int(box[1]), int(box[2]), int(box[3])
        img_copy = cv2.rectangle(
            img_copy,
            (x,y),(x+w, y+h),
            color, thickness)
        if draw_center:
            center = (x+w//2, y+h//2)
            img_copy = cv2.circle(img_copy, center=center, radius=3, color=(0,255,0), thickness=2)
    return img_copy


def read_data(annotations=Path(ANNOTATIONS_PATH)):
    """
        Reads annotations data from .csv file. Must contain columns: image_name, bbox_x, bbox_y, bbox_width, bbox_height.

        Arguments:
        annotations_path -- string or Path specifying path of annotations file

        Returns:
        data -- list of dictionaries containing path, number of boxes and boxes itself
    """
    data = []

    boxes = pd.read_csv(annotations)
    image_names = boxes['image_name'].unique()

    for image_name in image_names:
        cur_boxes = boxes[boxes['image_name'] == image_name]
        img_data = {
            'file_path': image_name,
            'box_nb': len(cur_boxes),
            'boxes': []}
        stamp_nb = img_data['box_nb']
        if stamp_nb <= STAMP_NB_MAX:
            img_data['boxes'] = cur_boxes[['bbox_x', 'bbox_y','bbox_width','bbox_height']].values
        data.append(img_data)
    return data


def boxes_to_tensor(boxes):
    """
        Convert list of boxes (and labels) to tensor format
        
        Arguments:
        boxes -- list of boxes

        Returns:
        boxes_tensor -- tensor of shape (S, S, BOX, 5)
    """
    boxes_tensor = torch.zeros((S, S, BOX, 5))
    cell_w, cell_h = W/S, H/S
    for i, box in enumerate(boxes):
        x, y, w, h = box
        # normalize xywh with cell_size
        x, y, w, h = x / cell_w, y / cell_h, w / cell_w, h / cell_h
        center_x, center_y = x + w / 2, y + h / 2
        grid_x = int(np.floor(center_x))
        grid_y = int(np.floor(center_y))
        
        if grid_x < S and grid_y < S:
            boxes_tensor[grid_y, grid_x, :, 0:4] = torch.tensor(BOX * [[center_x - grid_x, center_y - grid_y, w, h]])
            boxes_tensor[grid_y, grid_x, :, 4]  = torch.tensor(BOX * [1.])
    return boxes_tensor


def target_tensor_to_boxes(boxes_tensor, output_threshold=OUTPUT_THRESH):
    """
        Recover target tensor (tensor output of dataset) to bboxes.
        Arguments:
            boxes_tensor -- tensor of shape (S, S, BOX, 5)
        Returns:
            boxes -- list of boxes, each box is [x, y, w, h]
    """
    cell_w, cell_h = W/S, H/S
    boxes = []
    for i in range(S):
        for j in range(S):
            for b in range(BOX):
                data = boxes_tensor[i,j,b]
                x_center,y_center, w, h, obj_prob = data[0], data[1], data[2], data[3], data[4]
                if obj_prob > output_threshold:
                    x, y = x_center+j-w/2, y_center+i-h/2
                    x,y,w,h = x*cell_w, y*cell_h, w*cell_w, h*cell_h
                    box = [x,y,w,h]
                    boxes.append(box)
    return boxes    


def overlap(interval_1, interval_2):
    """
        Calculates length of overlap between two intervals.

        Arguments:
        interval_1 -- list or tuple of shape (2,) containing endpoints of the first interval
        interval_2 -- list or tuple of shape (2, 2) containing endpoints of the second interval

        Returns:
        overlap -- length of overlap
    """
    x1, x2 = interval_1
    x3, x4 = interval_2
    if x3 < x1:
        if x4 < x1:
            return 0
        else:
            return min(x2,x4) - x1
    else:
        if x2 < x3:
            return 0
        else:
            return min(x2,x4) - x3


def compute_iou(box1, box2):
    """
        Compute IOU between box1 and box2.

        Argmunets:
        box1 -- list of shape (5, ). Represents the first box
        box2 -- list of shape (5, ). Represents the second box
        Each box is [x, y, w, h, prob]

        Returns:
        iou -- intersection over union score between two boxes
    """
    x1,y1,w1,h1 = box1[0], box1[1], box1[2], box1[3]
    x2,y2,w2,h2 = box2[0], box2[1], box2[2], box2[3]

    area1, area2 = w1*h1, w2*h2
    intersect_w = overlap((x1,x1+w1), (x2,x2+w2))
    intersect_h = overlap((y1,y1+h1), (y2,y2+w2))
    if intersect_w == w1 and intersect_h == h1 or intersect_w == w2 and intersect_h == h2:
        return 1.
    intersect_area = intersect_w*intersect_h
    iou = intersect_area/(area1 + area2 - intersect_area)
    return iou


def nonmax_suppression(boxes, iou_thresh = IOU_THRESH):
    """
        Removes ovelap bboxes

        Arguments:
        boxes -- list of shape (None, 5)
        iou_thresh -- maximal value of iou when boxes are considered different
        Each box is [x, y, w, h, prob]

        Returns:
        boxes -- list of shape (None, 5) with removed overlapping boxes
    """
    boxes = sorted(boxes, key=lambda x: x[4], reverse=True)
    for i, current_box in enumerate(boxes):
        if current_box[4] <= 0:
            continue
        for j in range(i+1, len(boxes)):
            iou = compute_iou(current_box, boxes[j])
            if iou > iou_thresh:
                boxes[j][4] = 0
    boxes = [box for box in boxes if box[4] > 0]
    return boxes


    
def yolo_head(yolo_output):
    """
        Converts a yolo output tensor to separate tensors of coordinates, shapes and probabilities.

        Arguments:
        yolo_output -- tensor of shape (batch_size, S, S, BOX, 5)

        Returns:
        xy -- tensor of shape (batch_size, S, S, BOX, 2) containing coordinates of centers of found boxes for each anchor in each grid cell
        wh -- tensor of shape (batch_size, S, S, BOX, 2) containing width and height of found boxes for each anchor in each grid cell
        prob -- tensor of shape (batch_size, S, S, BOX, 1) containing the probability of presence of boxes for each anchor in each grid cell
    """
    xy = torch.sigmoid(yolo_output[..., 0:2])
    anchors_wh = torch.tensor(ANCHORS, device=yolo_output.device).view(1, 1, 1, len(ANCHORS), 2)
    wh = torch.exp(yolo_output[..., 2:4]) * anchors_wh
    prob = torch.sigmoid(yolo_output[..., 4:5])
    return xy, wh, prob

def process_target(target):
    xy = target[..., 0:2]
    wh = target[..., 2:4]
    prob = target[..., 4:5]
    return xy, wh, prob