287 lines
12 KiB
Python
287 lines
12 KiB
Python
|
import numpy as np
|
||
|
|
import cv2
|
||
|
|
from matplotlib import pyplot as plt
|
||
|
|
import torch
|
||
|
|
|
||
|
|
|
||
|
|
def data_visualizer(img, idx_to_class, path, bbox=None, pred=None):
|
||
|
|
"""
|
||
|
|
Data visualizer on the original image. Support both GT box input and proposal input.
|
||
|
|
|
||
|
|
Input:
|
||
|
|
- img: PIL Image input
|
||
|
|
- idx_to_class: Mapping from the index (0-19) to the class name
|
||
|
|
- bbox: GT bbox (in red, optional), a tensor of shape Nx5, where N is
|
||
|
|
the number of GT boxes, 5 indicates (x_tl, y_tl, x_br, y_br, class)
|
||
|
|
- pred: Predicted bbox (in green, optional), a tensor of shape N'x6, where
|
||
|
|
N' is the number of predicted boxes, 6 indicates
|
||
|
|
(x_tl, y_tl, x_br, y_br, class, object confidence score)
|
||
|
|
"""
|
||
|
|
|
||
|
|
img_copy = np.array(img).astype('uint8')
|
||
|
|
|
||
|
|
if bbox is not None:
|
||
|
|
for bbox_idx in range(bbox.shape[0]):
|
||
|
|
one_bbox = bbox[bbox_idx][:4].numpy().astype('int')
|
||
|
|
cv2.rectangle(img_copy, (one_bbox[0], one_bbox[1]), (one_bbox[2],
|
||
|
|
one_bbox[3]), (255, 0, 0), 2)
|
||
|
|
if bbox.shape[1] > 4: # if class info provided
|
||
|
|
obj_cls = idx_to_class[bbox[bbox_idx][4].item()]
|
||
|
|
cv2.putText(img_copy, '%s' % (obj_cls),
|
||
|
|
(one_bbox[0], one_bbox[1]+15),
|
||
|
|
cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 255), thickness=1)
|
||
|
|
|
||
|
|
if pred is not None:
|
||
|
|
for bbox_idx in range(pred.shape[0]):
|
||
|
|
one_bbox = pred[bbox_idx][:4].numpy().astype('int')
|
||
|
|
cv2.rectangle(img_copy, (one_bbox[0], one_bbox[1]), (one_bbox[2],
|
||
|
|
one_bbox[3]), (0, 255, 0), 2)
|
||
|
|
|
||
|
|
if pred.shape[1] > 4: # if class and conf score info provided
|
||
|
|
obj_cls = idx_to_class[pred[bbox_idx][4].item()]
|
||
|
|
conf_score = pred[bbox_idx][5].item()
|
||
|
|
cv2.putText(img_copy, '%s, %.2f' % (obj_cls, conf_score),
|
||
|
|
(one_bbox[0], one_bbox[1]+15),
|
||
|
|
cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 255), thickness=1)
|
||
|
|
|
||
|
|
plt.imshow(img_copy)
|
||
|
|
plt.axis('off')
|
||
|
|
plt.title(path)
|
||
|
|
plt.savefig(path)
|
||
|
|
plt.close()
|
||
|
|
|
||
|
|
|
||
|
|
def coord_trans(bbox, bbox_batch_idx, w_pixel, h_pixel, w_amap=7, h_amap=7, mode='a2p'):
|
||
|
|
"""
|
||
|
|
Coordinate transformation function. It converts the box coordinate from
|
||
|
|
the image coordinate system to the activation map coordinate system and vice versa.
|
||
|
|
In our case, the input image will have a few hundred of pixels in
|
||
|
|
width/height while the activation map is of size 7x7.
|
||
|
|
|
||
|
|
Input:
|
||
|
|
- bbox: Could be either bbox, anchor, or proposal, of shape Mx4
|
||
|
|
- bbox_batch_idx: Index of the image that each bbox belongs to, of shape M
|
||
|
|
- w_pixel: Number of pixels in the width side of the original image, of shape B
|
||
|
|
- h_pixel: Number of pixels in the height side of the original image, of shape B
|
||
|
|
- w_amap: Number of pixels in the width side of the activation map, scalar
|
||
|
|
- h_amap: Number of pixels in the height side of the activation map, scalar
|
||
|
|
- mode: Whether transfer from the original image to activation map ('p2a') or
|
||
|
|
the opposite ('a2p')
|
||
|
|
|
||
|
|
Output:
|
||
|
|
- resized_bbox: Resized box coordinates, of the same shape as the input bbox
|
||
|
|
"""
|
||
|
|
|
||
|
|
assert mode in ('p2a', 'a2p'), 'invalid coordinate transformation mode!'
|
||
|
|
assert bbox.shape[-1] >= 4, 'the transformation is applied to the first 4 values of dim -1'
|
||
|
|
|
||
|
|
if bbox.shape[0] == 0: # corner cases
|
||
|
|
return bbox
|
||
|
|
|
||
|
|
resized_bbox = bbox.clone()
|
||
|
|
|
||
|
|
if mode == 'p2a':
|
||
|
|
# pixel to activation
|
||
|
|
width_ratio = w_pixel[bbox_batch_idx] * 1. / w_amap
|
||
|
|
height_ratio = h_pixel[bbox_batch_idx] * 1. / h_amap
|
||
|
|
resized_bbox[:, [0, 2]] /= width_ratio.view(-1, 1)
|
||
|
|
resized_bbox[:, [1, 3]] /= height_ratio.view(-1, 1)
|
||
|
|
else:
|
||
|
|
# activation to pixel
|
||
|
|
width_ratio = w_pixel[bbox_batch_idx] * 1. / w_amap
|
||
|
|
height_ratio = h_pixel[bbox_batch_idx] * 1. / h_amap
|
||
|
|
resized_bbox[:, [0, 2]] *= width_ratio.view(-1, 1)
|
||
|
|
resized_bbox[:, [1, 3]] *= height_ratio.view(-1, 1)
|
||
|
|
|
||
|
|
return resized_bbox
|
||
|
|
|
||
|
|
|
||
|
|
def generate_anchor(anc_per_grid, grid):
|
||
|
|
"""
|
||
|
|
Anchor generator.
|
||
|
|
|
||
|
|
Inputs:
|
||
|
|
- anc_per_grid: Tensor of shape (A, 2) giving the shapes of anchor boxes to
|
||
|
|
consider at each point in the grid. anc_per_grid[a] = (w, h) gives the width
|
||
|
|
and height of the a'th anchor shape.
|
||
|
|
- grid: Tensor of shape (B, H', W', 2) giving the (x, y) coordinates of the
|
||
|
|
center of each feature from the backbone feature map. This is the tensor
|
||
|
|
returned from GenerateGrid.
|
||
|
|
|
||
|
|
Outputs:
|
||
|
|
- anchors: Tensor of shape (B, A, H', W', 4) giving the positions of all
|
||
|
|
anchor boxes for the entire image. anchors[b, a, h, w] is an anchor box
|
||
|
|
centered at grid[b, h, w], whose shape is given by anc[a]; we parameterize
|
||
|
|
boxes as anchors[b, a, h, w] = (x_tl, y_tl, x_br, y_br), where (x_tl, y_tl)
|
||
|
|
and (x_br, y_br) give the xy coordinates of the top-left and bottom-right
|
||
|
|
corners of the box.
|
||
|
|
"""
|
||
|
|
A, _ = anc_per_grid.shape
|
||
|
|
B, H, W, _ = grid.shape
|
||
|
|
anc_per_grid = anc_per_grid.to(grid)
|
||
|
|
|
||
|
|
anc_per_grid = anc_per_grid.view(1, A, 1, 1, -1).repeat(B, 1, H, W, 1)
|
||
|
|
grid = grid.view(B, 1, H, W, -1).repeat(1, A, 1, 1, 1)
|
||
|
|
|
||
|
|
x1y1 = grid - anc_per_grid / 2
|
||
|
|
x2y2 = grid + anc_per_grid / 2
|
||
|
|
anchors = torch.cat([x1y1, x2y2], dim=-1)
|
||
|
|
|
||
|
|
return anchors
|
||
|
|
|
||
|
|
|
||
|
|
def compute_iou(anchors, bboxes):
|
||
|
|
"""
|
||
|
|
Compute the intersection-over-union between anchors and gts.
|
||
|
|
|
||
|
|
Inputs:
|
||
|
|
- anchors: Anchor boxes, of shape (M, 4), where M is the number of proposals
|
||
|
|
- bboxes: GT boxes of shape (N, 4), where N is the number of GT boxes,
|
||
|
|
4 indicates (x_{lr}^{gt}, y_{lr}^{gt}, x_{rb}^{gt}, y_{rb}^{gt})
|
||
|
|
|
||
|
|
Outputs:
|
||
|
|
- iou: IoU matrix of shape (M, N)
|
||
|
|
"""
|
||
|
|
iou = None
|
||
|
|
##############################################################################
|
||
|
|
# TODO: Given anchors and gt bboxes, #
|
||
|
|
# compute the iou between each anchor and gt bbox. #
|
||
|
|
##############################################################################
|
||
|
|
pass
|
||
|
|
|
||
|
|
##############################################################################
|
||
|
|
# END OF YOUR CODE #
|
||
|
|
##############################################################################
|
||
|
|
|
||
|
|
return iou
|
||
|
|
|
||
|
|
|
||
|
|
def compute_offsets(anchors, bboxes):
|
||
|
|
"""
|
||
|
|
Compute the offsets between anchors and gts.
|
||
|
|
|
||
|
|
Inputs:
|
||
|
|
- anchors: Anchor boxes, of shape (M, 4)
|
||
|
|
- bboxes: GT boxes of shape (M, 4),
|
||
|
|
4 indicates (x_{lr}^{gt}, y_{lr}^{gt}, x_{rb}^{gt}, y_{rb}^{gt})
|
||
|
|
|
||
|
|
Outputs:
|
||
|
|
- offsets: offsets of shape (M, 4)
|
||
|
|
"""
|
||
|
|
wh_offsets = torch.log((bboxes[:, 2:4] - bboxes[:, :2]) \
|
||
|
|
/ (anchors[:, 2:4] - anchors[:, :2]))
|
||
|
|
|
||
|
|
xy_offsets = (bboxes[:, :2] + bboxes[:, 2:4] - \
|
||
|
|
anchors[:, :2] - anchors[:, 2:4]) / 2.
|
||
|
|
|
||
|
|
xy_offsets /= (anchors[:, 2:4] - anchors[:, :2])
|
||
|
|
|
||
|
|
offsets = torch.cat((xy_offsets, wh_offsets), dim=-1)
|
||
|
|
|
||
|
|
return offsets
|
||
|
|
|
||
|
|
|
||
|
|
def generate_proposal(anchors, offsets):
|
||
|
|
"""
|
||
|
|
Proposal generator.
|
||
|
|
|
||
|
|
Inputs:
|
||
|
|
- anchors: Anchor boxes, of shape (M, 4). Anchors are represented
|
||
|
|
by the coordinates of their top-left and bottom-right corners.
|
||
|
|
- offsets: Transformations of shape (M, 4) that will be used to
|
||
|
|
convert anchor boxes into region proposals. The transformation
|
||
|
|
offsets[m] = (tx, ty, tw, th) will be applied to the anchor
|
||
|
|
anchors[m].
|
||
|
|
|
||
|
|
Outputs:
|
||
|
|
- proposals: Region proposals of shape (M, 4), represented by the
|
||
|
|
coordinates of their top-left and bottom-right corners. Applying the
|
||
|
|
transform offsets[m] to the anchor[m] should give the
|
||
|
|
proposal proposals[m].
|
||
|
|
|
||
|
|
"""
|
||
|
|
proposals = None
|
||
|
|
##############################################################################
|
||
|
|
# TODO: Given anchor coordinates and the proposed offset for each anchor, #
|
||
|
|
# compute the proposal coordinates using the transformation formulas above. #
|
||
|
|
##############################################################################
|
||
|
|
# Replace "pass" statement with your code
|
||
|
|
pass
|
||
|
|
|
||
|
|
##############################################################################
|
||
|
|
# END OF YOUR CODE #
|
||
|
|
##############################################################################
|
||
|
|
|
||
|
|
return proposals
|
||
|
|
|
||
|
|
|
||
|
|
@torch.no_grad()
|
||
|
|
def assign_label(proposals, bboxes, background_id, pos_thresh=0.5, neg_thresh=0.5, pos_fraction=0.25):
|
||
|
|
"""
|
||
|
|
Determine the activated (positive) and negative proposals for model training.
|
||
|
|
|
||
|
|
For Fast R-CNN - Positive proposals are defined Any of the two
|
||
|
|
(i) the proposal/proposals with the highest IoU overlap with a GT box, or
|
||
|
|
(ii) a proposal that has an IoU overlap higher than positive threshold with any GT box.
|
||
|
|
Note: One proposal can match at most one GT box (the one with the largest IoU overlapping).
|
||
|
|
|
||
|
|
We assign a negative label to a proposal if its IoU ratio is lower than
|
||
|
|
a threshold value for all GT boxes. Proposals that are neither positive nor negative
|
||
|
|
do not contribute to the training objective.
|
||
|
|
|
||
|
|
Main steps include:
|
||
|
|
i) Decide activated and negative proposals based on the IoU matrix.
|
||
|
|
ii) Compute GT confidence score/offsets/object class on the positive proposals.
|
||
|
|
iii) Compute GT confidence score on the negative proposals.
|
||
|
|
|
||
|
|
Inputs:
|
||
|
|
- proposal: Proposal boxes, of shape (M, 4), where M is the number of proposals
|
||
|
|
- bboxes: GT boxes of shape Nx5, where N is the number of GT boxes,
|
||
|
|
5 indicates (x_{lr}^{gt}, y_{lr}^{gt}, x_{rb}^{gt}, y_{rb}^{gt}) and class index
|
||
|
|
- background_id: Class id of the background class
|
||
|
|
- pos_thresh: Positive threshold value
|
||
|
|
- neg_thresh: Negative threshold value
|
||
|
|
- pos_fraction: a factor balancing pos/neg proposals
|
||
|
|
|
||
|
|
Outputs:
|
||
|
|
- activated_anc_mask: a binary mask indicating the activated proposals, of shape M
|
||
|
|
- negative_anc_mask: a binary mask indicating the negative proposals, of shape M
|
||
|
|
- GT_class: GT class category on all proposals, background class for non-activated proposals,
|
||
|
|
of shape M
|
||
|
|
- bboxes: GT bboxes on activated proposals, of shape M'x4, where M' is the number of
|
||
|
|
activated proposals
|
||
|
|
"""
|
||
|
|
M = proposals.shape[0]
|
||
|
|
N = bboxes.shape[0]
|
||
|
|
iou_mat = compute_iou(proposals, bboxes[:, :4])
|
||
|
|
|
||
|
|
# activated/positive proposals
|
||
|
|
max_iou_per_anc, max_iou_per_anc_ind = iou_mat.max(dim=-1)
|
||
|
|
max_iou_per_box = iou_mat.max(dim=0, keepdim=True)[0]
|
||
|
|
activated_anc_mask = (iou_mat == max_iou_per_box) & (max_iou_per_box > 0)
|
||
|
|
activated_anc_mask |= (iou_mat > pos_thresh) # using the pos_thresh condition as well
|
||
|
|
activated_anc_mask = activated_anc_mask.max(dim=-1)[0] # (M, )
|
||
|
|
activated_anc_ind = torch.nonzero(activated_anc_mask.view(-1)).squeeze(-1)
|
||
|
|
|
||
|
|
# GT class
|
||
|
|
box_cls = bboxes[:, 4].long().view(1, N).expand(M, N)
|
||
|
|
# if a proposal matches multiple GT boxes, choose the box with the largest iou
|
||
|
|
GT_class = torch.gather(box_cls, -1, max_iou_per_anc_ind.unsqueeze(-1)).squeeze(-1) # M
|
||
|
|
GT_class[~activated_anc_mask] = background_id
|
||
|
|
|
||
|
|
# GT bboxes
|
||
|
|
bboxes_expand = bboxes[:, :4].view(1, N, 4).expand((M, N, 4))
|
||
|
|
bboxes = torch.gather(bboxes_expand, -2, max_iou_per_anc_ind.unsqueeze(-1) \
|
||
|
|
.unsqueeze(-1).expand(M, 1, 4)).view(M, 4)
|
||
|
|
bboxes = bboxes[activated_anc_ind]
|
||
|
|
|
||
|
|
# negative anchors
|
||
|
|
negative_anc_mask = (max_iou_per_anc < neg_thresh)
|
||
|
|
negative_anc_ind = torch.nonzero(negative_anc_mask.view(-1)).squeeze(-1)
|
||
|
|
# balance pos/neg anchors, random choose
|
||
|
|
num_neg = int(activated_anc_ind.shape[0] * (1 - pos_fraction) / pos_fraction)
|
||
|
|
negative_anc_ind = negative_anc_ind[torch.randint(0, negative_anc_ind.shape[0], (num_neg,))]
|
||
|
|
negative_anc_mask = torch.zeros_like(negative_anc_mask)
|
||
|
|
negative_anc_mask[negative_anc_ind] = 1
|
||
|
|
|
||
|
|
return activated_anc_mask, negative_anc_mask, GT_class, bboxes
|