1920 lines
72 KiB
Python
1920 lines
72 KiB
Python
"""
|
|
coding=utf-8
|
|
Copyright 2018, Antonio Mendoza Hao Tan, Mohit Bansal
|
|
Adapted From Facebook Inc, Detectron2 && Huggingface Co.
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.import copy
|
|
"""
|
|
import itertools
|
|
import math
|
|
import os
|
|
from abc import ABCMeta, abstractmethod
|
|
from collections import OrderedDict, namedtuple
|
|
from typing import Dict, List, Tuple
|
|
|
|
import numpy as np
|
|
import torch
|
|
from torch import nn
|
|
from torch.nn.modules.batchnorm import BatchNorm2d
|
|
from torchvision.ops import RoIPool
|
|
from torchvision.ops.boxes import batched_nms, nms
|
|
|
|
from utils import WEIGHTS_NAME, Config, cached_path, hf_bucket_url, is_remote_url, load_checkpoint
|
|
|
|
|
|
# other:
|
|
def norm_box(boxes, raw_sizes):
|
|
if not isinstance(boxes, torch.Tensor):
|
|
normalized_boxes = boxes.copy()
|
|
else:
|
|
normalized_boxes = boxes.clone()
|
|
normalized_boxes[:, :, (0, 2)] /= raw_sizes[:, 1]
|
|
normalized_boxes[:, :, (1, 3)] /= raw_sizes[:, 0]
|
|
return normalized_boxes
|
|
|
|
|
|
def pad_list_tensors(
|
|
list_tensors,
|
|
preds_per_image,
|
|
max_detections=None,
|
|
return_tensors=None,
|
|
padding=None,
|
|
pad_value=0,
|
|
location=None,
|
|
):
|
|
"""
|
|
location will always be cpu for np tensors
|
|
"""
|
|
if location is None:
|
|
location = "cpu"
|
|
assert return_tensors in {"pt", "np", None}
|
|
assert padding in {"max_detections", "max_batch", None}
|
|
new = []
|
|
if padding is None:
|
|
if return_tensors is None:
|
|
return list_tensors
|
|
elif return_tensors == "pt":
|
|
if not isinstance(list_tensors, torch.Tensor):
|
|
return torch.stack(list_tensors).to(location)
|
|
else:
|
|
return list_tensors.to(location)
|
|
else:
|
|
if not isinstance(list_tensors, list):
|
|
return np.array(list_tensors.to(location))
|
|
else:
|
|
return list_tensors.to(location)
|
|
if padding == "max_detections":
|
|
assert max_detections is not None, "specify max number of detections per batch"
|
|
elif padding == "max_batch":
|
|
max_detections = max(preds_per_image)
|
|
for i in range(len(list_tensors)):
|
|
too_small = False
|
|
tensor_i = list_tensors.pop(0)
|
|
if tensor_i.ndim < 2:
|
|
too_small = True
|
|
tensor_i = tensor_i.unsqueeze(-1)
|
|
assert isinstance(tensor_i, torch.Tensor)
|
|
tensor_i = nn.functional.pad(
|
|
input=tensor_i,
|
|
pad=(0, 0, 0, max_detections - preds_per_image[i]),
|
|
mode="constant",
|
|
value=pad_value,
|
|
)
|
|
if too_small:
|
|
tensor_i = tensor_i.squeeze(-1)
|
|
if return_tensors is None:
|
|
if location == "cpu":
|
|
tensor_i = tensor_i.cpu()
|
|
tensor_i = tensor_i.tolist()
|
|
if return_tensors == "np":
|
|
if location == "cpu":
|
|
tensor_i = tensor_i.cpu()
|
|
tensor_i = tensor_i.numpy()
|
|
else:
|
|
if location == "cpu":
|
|
tensor_i = tensor_i.cpu()
|
|
new.append(tensor_i)
|
|
if return_tensors == "np":
|
|
return np.stack(new, axis=0)
|
|
elif return_tensors == "pt" and not isinstance(new, torch.Tensor):
|
|
return torch.stack(new, dim=0)
|
|
else:
|
|
return list_tensors
|
|
|
|
|
|
def do_nms(boxes, scores, image_shape, score_thresh, nms_thresh, mind, maxd):
|
|
scores = scores[:, :-1]
|
|
num_bbox_reg_classes = boxes.shape[1] // 4
|
|
# Convert to Boxes to use the `clip` function ...
|
|
boxes = boxes.reshape(-1, 4)
|
|
_clip_box(boxes, image_shape)
|
|
boxes = boxes.view(-1, num_bbox_reg_classes, 4) # R x C x 4
|
|
|
|
# Select max scores
|
|
max_scores, max_classes = scores.max(1) # R x C --> R
|
|
num_objs = boxes.size(0)
|
|
boxes = boxes.view(-1, 4)
|
|
idxs = torch.arange(num_objs).to(boxes.device) * num_bbox_reg_classes + max_classes
|
|
max_boxes = boxes[idxs] # Select max boxes according to the max scores.
|
|
|
|
# Apply NMS
|
|
keep = nms(max_boxes, max_scores, nms_thresh)
|
|
keep = keep[:maxd]
|
|
if keep.shape[-1] >= mind and keep.shape[-1] <= maxd:
|
|
max_boxes, max_scores = max_boxes[keep], max_scores[keep]
|
|
classes = max_classes[keep]
|
|
return max_boxes, max_scores, classes, keep
|
|
else:
|
|
return None
|
|
|
|
|
|
# Helper Functions
|
|
def _clip_box(tensor, box_size: Tuple[int, int]):
|
|
assert torch.isfinite(tensor).all(), "Box tensor contains infinite or NaN!"
|
|
h, w = box_size
|
|
tensor[:, 0].clamp_(min=0, max=w)
|
|
tensor[:, 1].clamp_(min=0, max=h)
|
|
tensor[:, 2].clamp_(min=0, max=w)
|
|
tensor[:, 3].clamp_(min=0, max=h)
|
|
|
|
|
|
def _nonempty_boxes(box, threshold: float = 0.0) -> torch.Tensor:
|
|
widths = box[:, 2] - box[:, 0]
|
|
heights = box[:, 3] - box[:, 1]
|
|
keep = (widths > threshold) & (heights > threshold)
|
|
return keep
|
|
|
|
|
|
def get_norm(norm, out_channels):
|
|
if isinstance(norm, str):
|
|
if len(norm) == 0:
|
|
return None
|
|
norm = {
|
|
"BN": BatchNorm2d,
|
|
"GN": lambda channels: nn.GroupNorm(32, channels),
|
|
"nnSyncBN": nn.SyncBatchNorm, # keep for debugging
|
|
"": lambda x: x,
|
|
}[norm]
|
|
return norm(out_channels)
|
|
|
|
|
|
def _create_grid_offsets(size: List[int], stride: int, offset: float, device):
|
|
grid_height, grid_width = size
|
|
shifts_x = torch.arange(
|
|
offset * stride,
|
|
grid_width * stride,
|
|
step=stride,
|
|
dtype=torch.float32,
|
|
device=device,
|
|
)
|
|
shifts_y = torch.arange(
|
|
offset * stride,
|
|
grid_height * stride,
|
|
step=stride,
|
|
dtype=torch.float32,
|
|
device=device,
|
|
)
|
|
|
|
shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x)
|
|
shift_x = shift_x.reshape(-1)
|
|
shift_y = shift_y.reshape(-1)
|
|
return shift_x, shift_y
|
|
|
|
|
|
def build_backbone(cfg):
|
|
input_shape = ShapeSpec(channels=len(cfg.MODEL.PIXEL_MEAN))
|
|
norm = cfg.RESNETS.NORM
|
|
stem = BasicStem(
|
|
in_channels=input_shape.channels,
|
|
out_channels=cfg.RESNETS.STEM_OUT_CHANNELS,
|
|
norm=norm,
|
|
caffe_maxpool=cfg.MODEL.MAX_POOL,
|
|
)
|
|
freeze_at = cfg.BACKBONE.FREEZE_AT
|
|
|
|
if freeze_at >= 1:
|
|
for p in stem.parameters():
|
|
p.requires_grad = False
|
|
|
|
out_features = cfg.RESNETS.OUT_FEATURES
|
|
depth = cfg.RESNETS.DEPTH
|
|
num_groups = cfg.RESNETS.NUM_GROUPS
|
|
width_per_group = cfg.RESNETS.WIDTH_PER_GROUP
|
|
bottleneck_channels = num_groups * width_per_group
|
|
in_channels = cfg.RESNETS.STEM_OUT_CHANNELS
|
|
out_channels = cfg.RESNETS.RES2_OUT_CHANNELS
|
|
stride_in_1x1 = cfg.RESNETS.STRIDE_IN_1X1
|
|
res5_dilation = cfg.RESNETS.RES5_DILATION
|
|
assert res5_dilation in {1, 2}, "res5_dilation cannot be {}.".format(res5_dilation)
|
|
|
|
num_blocks_per_stage = {50: [3, 4, 6, 3], 101: [3, 4, 23, 3], 152: [3, 8, 36, 3]}[depth]
|
|
|
|
stages = []
|
|
out_stage_idx = [{"res2": 2, "res3": 3, "res4": 4, "res5": 5}[f] for f in out_features]
|
|
max_stage_idx = max(out_stage_idx)
|
|
for idx, stage_idx in enumerate(range(2, max_stage_idx + 1)):
|
|
dilation = res5_dilation if stage_idx == 5 else 1
|
|
first_stride = 1 if idx == 0 or (stage_idx == 5 and dilation == 2) else 2
|
|
stage_kargs = {
|
|
"num_blocks": num_blocks_per_stage[idx],
|
|
"first_stride": first_stride,
|
|
"in_channels": in_channels,
|
|
"bottleneck_channels": bottleneck_channels,
|
|
"out_channels": out_channels,
|
|
"num_groups": num_groups,
|
|
"norm": norm,
|
|
"stride_in_1x1": stride_in_1x1,
|
|
"dilation": dilation,
|
|
}
|
|
|
|
stage_kargs["block_class"] = BottleneckBlock
|
|
blocks = ResNet.make_stage(**stage_kargs)
|
|
in_channels = out_channels
|
|
out_channels *= 2
|
|
bottleneck_channels *= 2
|
|
|
|
if freeze_at >= stage_idx:
|
|
for block in blocks:
|
|
block.freeze()
|
|
stages.append(blocks)
|
|
|
|
return ResNet(stem, stages, out_features=out_features)
|
|
|
|
|
|
def find_top_rpn_proposals(
|
|
proposals,
|
|
pred_objectness_logits,
|
|
images,
|
|
image_sizes,
|
|
nms_thresh,
|
|
pre_nms_topk,
|
|
post_nms_topk,
|
|
min_box_side_len,
|
|
training,
|
|
):
|
|
"""Args:
|
|
proposals (list[Tensor]): (L, N, Hi*Wi*A, 4).
|
|
pred_objectness_logits: tensors of length L.
|
|
nms_thresh (float): IoU threshold to use for NMS
|
|
pre_nms_topk (int): before nms
|
|
post_nms_topk (int): after nms
|
|
min_box_side_len (float): minimum proposal box side
|
|
training (bool): True if proposals are to be used in training,
|
|
Returns:
|
|
results (List[Dict]): stores post_nms_topk object proposals for image i.
|
|
"""
|
|
num_images = len(images)
|
|
device = proposals[0].device
|
|
|
|
# 1. Select top-k anchor for every level and every image
|
|
topk_scores = [] # #lvl Tensor, each of shape N x topk
|
|
topk_proposals = []
|
|
level_ids = [] # #lvl Tensor, each of shape (topk,)
|
|
batch_idx = torch.arange(num_images, device=device)
|
|
for level_id, proposals_i, logits_i in zip(itertools.count(), proposals, pred_objectness_logits):
|
|
Hi_Wi_A = logits_i.shape[1]
|
|
num_proposals_i = min(pre_nms_topk, Hi_Wi_A)
|
|
|
|
# sort is faster than topk (https://github.com/pytorch/pytorch/issues/22812)
|
|
# topk_scores_i, topk_idx = logits_i.topk(num_proposals_i, dim=1)
|
|
logits_i, idx = logits_i.sort(descending=True, dim=1)
|
|
topk_scores_i = logits_i[batch_idx, :num_proposals_i]
|
|
topk_idx = idx[batch_idx, :num_proposals_i]
|
|
|
|
# each is N x topk
|
|
topk_proposals_i = proposals_i[batch_idx[:, None], topk_idx] # N x topk x 4
|
|
|
|
topk_proposals.append(topk_proposals_i)
|
|
topk_scores.append(topk_scores_i)
|
|
level_ids.append(torch.full((num_proposals_i,), level_id, dtype=torch.int64, device=device))
|
|
|
|
# 2. Concat all levels together
|
|
topk_scores = torch.cat(topk_scores, dim=1)
|
|
topk_proposals = torch.cat(topk_proposals, dim=1)
|
|
level_ids = torch.cat(level_ids, dim=0)
|
|
|
|
# if I change to batched_nms, I wonder if this will make a difference
|
|
# 3. For each image, run a per-level NMS, and choose topk results.
|
|
results = []
|
|
for n, image_size in enumerate(image_sizes):
|
|
boxes = topk_proposals[n]
|
|
scores_per_img = topk_scores[n]
|
|
# I will have to take a look at the boxes clip method
|
|
_clip_box(boxes, image_size)
|
|
# filter empty boxes
|
|
keep = _nonempty_boxes(boxes, threshold=min_box_side_len)
|
|
lvl = level_ids
|
|
if keep.sum().item() != len(boxes):
|
|
boxes, scores_per_img, lvl = (
|
|
boxes[keep],
|
|
scores_per_img[keep],
|
|
level_ids[keep],
|
|
)
|
|
|
|
keep = batched_nms(boxes, scores_per_img, lvl, nms_thresh)
|
|
keep = keep[:post_nms_topk]
|
|
|
|
res = (boxes[keep], scores_per_img[keep])
|
|
results.append(res)
|
|
|
|
# I wonder if it would be possible for me to pad all these things.
|
|
return results
|
|
|
|
|
|
def subsample_labels(labels, num_samples, positive_fraction, bg_label):
|
|
"""
|
|
Returns:
|
|
pos_idx, neg_idx (Tensor):
|
|
1D vector of indices. The total length of both is `num_samples` or fewer.
|
|
"""
|
|
positive = torch.nonzero((labels != -1) & (labels != bg_label)).squeeze(1)
|
|
negative = torch.nonzero(labels == bg_label).squeeze(1)
|
|
|
|
num_pos = int(num_samples * positive_fraction)
|
|
# protect against not enough positive examples
|
|
num_pos = min(positive.numel(), num_pos)
|
|
num_neg = num_samples - num_pos
|
|
# protect against not enough negative examples
|
|
num_neg = min(negative.numel(), num_neg)
|
|
|
|
# randomly select positive and negative examples
|
|
perm1 = torch.randperm(positive.numel(), device=positive.device)[:num_pos]
|
|
perm2 = torch.randperm(negative.numel(), device=negative.device)[:num_neg]
|
|
|
|
pos_idx = positive[perm1]
|
|
neg_idx = negative[perm2]
|
|
return pos_idx, neg_idx
|
|
|
|
|
|
def add_ground_truth_to_proposals(gt_boxes, proposals):
|
|
raise NotImplementedError()
|
|
|
|
|
|
def add_ground_truth_to_proposals_single_image(gt_boxes, proposals):
|
|
raise NotImplementedError()
|
|
|
|
|
|
def _fmt_box_list(box_tensor, batch_index: int):
|
|
repeated_index = torch.full(
|
|
(len(box_tensor), 1),
|
|
batch_index,
|
|
dtype=box_tensor.dtype,
|
|
device=box_tensor.device,
|
|
)
|
|
return torch.cat((repeated_index, box_tensor), dim=1)
|
|
|
|
|
|
def convert_boxes_to_pooler_format(box_lists: List[torch.Tensor]):
|
|
pooler_fmt_boxes = torch.cat(
|
|
[_fmt_box_list(box_list, i) for i, box_list in enumerate(box_lists)],
|
|
dim=0,
|
|
)
|
|
return pooler_fmt_boxes
|
|
|
|
|
|
def assign_boxes_to_levels(
|
|
box_lists: List[torch.Tensor],
|
|
min_level: int,
|
|
max_level: int,
|
|
canonical_box_size: int,
|
|
canonical_level: int,
|
|
):
|
|
box_sizes = torch.sqrt(torch.cat([boxes.area() for boxes in box_lists]))
|
|
# Eqn.(1) in FPN paper
|
|
level_assignments = torch.floor(canonical_level + torch.log2(box_sizes / canonical_box_size + 1e-8))
|
|
# clamp level to (min, max), in case the box size is too large or too small
|
|
# for the available feature maps
|
|
level_assignments = torch.clamp(level_assignments, min=min_level, max=max_level)
|
|
return level_assignments.to(torch.int64) - min_level
|
|
|
|
|
|
# Helper Classes
|
|
class _NewEmptyTensorOp(torch.autograd.Function):
|
|
@staticmethod
|
|
def forward(ctx, x, new_shape):
|
|
ctx.shape = x.shape
|
|
return x.new_empty(new_shape)
|
|
|
|
@staticmethod
|
|
def backward(ctx, grad):
|
|
shape = ctx.shape
|
|
return _NewEmptyTensorOp.apply(grad, shape), None
|
|
|
|
|
|
class ShapeSpec(namedtuple("_ShapeSpec", ["channels", "height", "width", "stride"])):
|
|
def __new__(cls, *, channels=None, height=None, width=None, stride=None):
|
|
return super().__new__(cls, channels, height, width, stride)
|
|
|
|
|
|
class Box2BoxTransform(object):
|
|
"""
|
|
This R-CNN transformation scales the box's width and height
|
|
by exp(dw), exp(dh) and shifts a box's center by the offset
|
|
(dx * width, dy * height).
|
|
"""
|
|
|
|
def __init__(self, weights: Tuple[float, float, float, float], scale_clamp: float = None):
|
|
"""
|
|
Args:
|
|
weights (4-element tuple): Scaling factors that are applied to the
|
|
(dx, dy, dw, dh) deltas. In Fast R-CNN, these were originally set
|
|
such that the deltas have unit variance; now they are treated as
|
|
hyperparameters of the system.
|
|
scale_clamp (float): When predicting deltas, the predicted box scaling
|
|
factors (dw and dh) are clamped such that they are <= scale_clamp.
|
|
"""
|
|
self.weights = weights
|
|
if scale_clamp is not None:
|
|
self.scale_clamp = scale_clamp
|
|
else:
|
|
"""
|
|
Value for clamping large dw and dh predictions.
|
|
The heuristic is that we clamp such that dw and dh are no larger
|
|
than what would transform a 16px box into a 1000px box
|
|
(based on a small anchor, 16px, and a typical image size, 1000px).
|
|
"""
|
|
self.scale_clamp = math.log(1000.0 / 16)
|
|
|
|
def get_deltas(self, src_boxes, target_boxes):
|
|
"""
|
|
Get box regression transformation deltas (dx, dy, dw, dh) that can be used
|
|
to transform the `src_boxes` into the `target_boxes`. That is, the relation
|
|
``target_boxes == self.apply_deltas(deltas, src_boxes)`` is true (unless
|
|
any delta is too large and is clamped).
|
|
Args:
|
|
src_boxes (Tensor): source boxes, e.g., object proposals
|
|
target_boxes (Tensor): target of the transformation, e.g., ground-truth
|
|
boxes.
|
|
"""
|
|
assert isinstance(src_boxes, torch.Tensor), type(src_boxes)
|
|
assert isinstance(target_boxes, torch.Tensor), type(target_boxes)
|
|
|
|
src_widths = src_boxes[:, 2] - src_boxes[:, 0]
|
|
src_heights = src_boxes[:, 3] - src_boxes[:, 1]
|
|
src_ctr_x = src_boxes[:, 0] + 0.5 * src_widths
|
|
src_ctr_y = src_boxes[:, 1] + 0.5 * src_heights
|
|
|
|
target_widths = target_boxes[:, 2] - target_boxes[:, 0]
|
|
target_heights = target_boxes[:, 3] - target_boxes[:, 1]
|
|
target_ctr_x = target_boxes[:, 0] + 0.5 * target_widths
|
|
target_ctr_y = target_boxes[:, 1] + 0.5 * target_heights
|
|
|
|
wx, wy, ww, wh = self.weights
|
|
dx = wx * (target_ctr_x - src_ctr_x) / src_widths
|
|
dy = wy * (target_ctr_y - src_ctr_y) / src_heights
|
|
dw = ww * torch.log(target_widths / src_widths)
|
|
dh = wh * torch.log(target_heights / src_heights)
|
|
|
|
deltas = torch.stack((dx, dy, dw, dh), dim=1)
|
|
assert (src_widths > 0).all().item(), "Input boxes to Box2BoxTransform are not valid!"
|
|
return deltas
|
|
|
|
def apply_deltas(self, deltas, boxes):
|
|
"""
|
|
Apply transformation `deltas` (dx, dy, dw, dh) to `boxes`.
|
|
Args:
|
|
deltas (Tensor): transformation deltas of shape (N, k*4), where k >= 1.
|
|
deltas[i] represents k potentially different class-specific
|
|
box transformations for the single box boxes[i].
|
|
boxes (Tensor): boxes to transform, of shape (N, 4)
|
|
"""
|
|
boxes = boxes.to(deltas.dtype)
|
|
|
|
widths = boxes[:, 2] - boxes[:, 0]
|
|
heights = boxes[:, 3] - boxes[:, 1]
|
|
ctr_x = boxes[:, 0] + 0.5 * widths
|
|
ctr_y = boxes[:, 1] + 0.5 * heights
|
|
|
|
wx, wy, ww, wh = self.weights
|
|
dx = deltas[:, 0::4] / wx
|
|
dy = deltas[:, 1::4] / wy
|
|
dw = deltas[:, 2::4] / ww
|
|
dh = deltas[:, 3::4] / wh
|
|
|
|
# Prevent sending too large values into torch.exp()
|
|
dw = torch.clamp(dw, max=self.scale_clamp)
|
|
dh = torch.clamp(dh, max=self.scale_clamp)
|
|
|
|
pred_ctr_x = dx * widths[:, None] + ctr_x[:, None]
|
|
pred_ctr_y = dy * heights[:, None] + ctr_y[:, None]
|
|
pred_w = torch.exp(dw) * widths[:, None]
|
|
pred_h = torch.exp(dh) * heights[:, None]
|
|
|
|
pred_boxes = torch.zeros_like(deltas)
|
|
pred_boxes[:, 0::4] = pred_ctr_x - 0.5 * pred_w # x1
|
|
pred_boxes[:, 1::4] = pred_ctr_y - 0.5 * pred_h # y1
|
|
pred_boxes[:, 2::4] = pred_ctr_x + 0.5 * pred_w # x2
|
|
pred_boxes[:, 3::4] = pred_ctr_y + 0.5 * pred_h # y2
|
|
return pred_boxes
|
|
|
|
|
|
class Matcher(object):
|
|
"""
|
|
This class assigns to each predicted "element" (e.g., a box) a ground-truth
|
|
element. Each predicted element will have exactly zero or one matches; each
|
|
ground-truth element may be matched to zero or more predicted elements.
|
|
The matching is determined by the MxN match_quality_matrix, that characterizes
|
|
how well each (ground-truth, prediction)-pair match each other. For example,
|
|
if the elements are boxes, this matrix may contain box intersection-over-union
|
|
overlap values.
|
|
The matcher returns (a) a vector of length N containing the index of the
|
|
ground-truth element m in [0, M) that matches to prediction n in [0, N).
|
|
(b) a vector of length N containing the labels for each prediction.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
thresholds: List[float],
|
|
labels: List[int],
|
|
allow_low_quality_matches: bool = False,
|
|
):
|
|
"""
|
|
Args:
|
|
thresholds (list): a list of thresholds used to stratify predictions
|
|
into levels.
|
|
labels (list): a list of values to label predictions belonging at
|
|
each level. A label can be one of {-1, 0, 1} signifying
|
|
{ignore, negative class, positive class}, respectively.
|
|
allow_low_quality_matches (bool): if True, produce additional matches or predictions with maximum match quality lower than high_threshold.
|
|
For example, thresholds = [0.3, 0.5] labels = [0, -1, 1] All predictions with iou < 0.3 will be marked with 0 and
|
|
thus will be considered as false positives while training. All predictions with 0.3 <= iou < 0.5 will be marked with -1 and
|
|
thus will be ignored. All predictions with 0.5 <= iou will be marked with 1 and thus will be considered as true positives.
|
|
"""
|
|
thresholds = thresholds[:]
|
|
assert thresholds[0] > 0
|
|
thresholds.insert(0, -float("inf"))
|
|
thresholds.append(float("inf"))
|
|
assert all(low <= high for (low, high) in zip(thresholds[:-1], thresholds[1:]))
|
|
assert all(label_i in [-1, 0, 1] for label_i in labels)
|
|
assert len(labels) == len(thresholds) - 1
|
|
self.thresholds = thresholds
|
|
self.labels = labels
|
|
self.allow_low_quality_matches = allow_low_quality_matches
|
|
|
|
def __call__(self, match_quality_matrix):
|
|
"""
|
|
Args:
|
|
match_quality_matrix (Tensor[float]): an MxN tensor, containing the pairwise quality between M ground-truth elements and N predicted
|
|
elements. All elements must be >= 0 (due to the us of `torch.nonzero` for selecting indices in :meth:`set_low_quality_matches_`).
|
|
Returns:
|
|
matches (Tensor[int64]): a vector of length N, where matches[i] is a matched ground-truth index in [0, M)
|
|
match_labels (Tensor[int8]): a vector of length N, where pred_labels[i] indicates true or false positive or ignored
|
|
"""
|
|
assert match_quality_matrix.dim() == 2
|
|
if match_quality_matrix.numel() == 0:
|
|
default_matches = match_quality_matrix.new_full((match_quality_matrix.size(1),), 0, dtype=torch.int64)
|
|
# When no gt boxes exist, we define IOU = 0 and therefore set labels
|
|
# to `self.labels[0]`, which usually defaults to background class 0
|
|
# To choose to ignore instead,
|
|
# can make labels=[-1,0,-1,1] + set appropriate thresholds
|
|
default_match_labels = match_quality_matrix.new_full(
|
|
(match_quality_matrix.size(1),), self.labels[0], dtype=torch.int8
|
|
)
|
|
return default_matches, default_match_labels
|
|
|
|
assert torch.all(match_quality_matrix >= 0)
|
|
|
|
# match_quality_matrix is M (gt) x N (predicted)
|
|
# Max over gt elements (dim 0) to find best gt candidate for each prediction
|
|
matched_vals, matches = match_quality_matrix.max(dim=0)
|
|
|
|
match_labels = matches.new_full(matches.size(), 1, dtype=torch.int8)
|
|
|
|
for l, low, high in zip(self.labels, self.thresholds[:-1], self.thresholds[1:]):
|
|
low_high = (matched_vals >= low) & (matched_vals < high)
|
|
match_labels[low_high] = l
|
|
|
|
if self.allow_low_quality_matches:
|
|
self.set_low_quality_matches_(match_labels, match_quality_matrix)
|
|
|
|
return matches, match_labels
|
|
|
|
def set_low_quality_matches_(self, match_labels, match_quality_matrix):
|
|
"""
|
|
Produce additional matches for predictions that have only low-quality matches.
|
|
Specifically, for each ground-truth G find the set of predictions that have
|
|
maximum overlap with it (including ties); for each prediction in that set, if
|
|
it is unmatched, then match it to the ground-truth G.
|
|
This function implements the RPN assignment case (i)
|
|
in Sec. 3.1.2 of Faster R-CNN.
|
|
"""
|
|
# For each gt, find the prediction with which it has highest quality
|
|
highest_quality_foreach_gt, _ = match_quality_matrix.max(dim=1)
|
|
# Find the highest quality match available, even if it is low, including ties.
|
|
# Note that the matches qualities must be positive due to the use of
|
|
# `torch.nonzero`.
|
|
of_quality_inds = match_quality_matrix == highest_quality_foreach_gt[:, None]
|
|
if of_quality_inds.dim() == 0:
|
|
(_, pred_inds_with_highest_quality) = of_quality_inds.unsqueeze(0).nonzero().unbind(1)
|
|
else:
|
|
(_, pred_inds_with_highest_quality) = of_quality_inds.nonzero().unbind(1)
|
|
match_labels[pred_inds_with_highest_quality] = 1
|
|
|
|
|
|
class RPNOutputs(object):
|
|
def __init__(
|
|
self,
|
|
box2box_transform,
|
|
anchor_matcher,
|
|
batch_size_per_image,
|
|
positive_fraction,
|
|
images,
|
|
pred_objectness_logits,
|
|
pred_anchor_deltas,
|
|
anchors,
|
|
boundary_threshold=0,
|
|
gt_boxes=None,
|
|
smooth_l1_beta=0.0,
|
|
):
|
|
"""
|
|
Args:
|
|
box2box_transform (Box2BoxTransform): :class:`Box2BoxTransform` instance for anchor-proposal transformations.
|
|
anchor_matcher (Matcher): :class:`Matcher` instance for matching anchors to ground-truth boxes; used to determine training labels.
|
|
batch_size_per_image (int): number of proposals to sample when training
|
|
positive_fraction (float): target fraction of sampled proposals that should be positive
|
|
images (ImageList): :class:`ImageList` instance representing N input images
|
|
pred_objectness_logits (list[Tensor]): A list of L elements. Element i is a tensor of shape (N, A, Hi, W)
|
|
pred_anchor_deltas (list[Tensor]): A list of L elements. Element i is a tensor of shape (N, A*4, Hi, Wi)
|
|
anchors (list[torch.Tensor]): nested list of boxes. anchors[i][j] at (n, l) stores anchor array for feature map l
|
|
boundary_threshold (int): if >= 0, then anchors that extend beyond the image boundary by more than boundary_thresh are not used in training.
|
|
gt_boxes (list[Boxes], optional): A list of N elements.
|
|
smooth_l1_beta (float): The transition point between L1 and L2 lossn. When set to 0, the loss becomes L1. When +inf, it is ignored
|
|
"""
|
|
self.box2box_transform = box2box_transform
|
|
self.anchor_matcher = anchor_matcher
|
|
self.batch_size_per_image = batch_size_per_image
|
|
self.positive_fraction = positive_fraction
|
|
self.pred_objectness_logits = pred_objectness_logits
|
|
self.pred_anchor_deltas = pred_anchor_deltas
|
|
|
|
self.anchors = anchors
|
|
self.gt_boxes = gt_boxes
|
|
self.num_feature_maps = len(pred_objectness_logits)
|
|
self.num_images = len(images)
|
|
self.boundary_threshold = boundary_threshold
|
|
self.smooth_l1_beta = smooth_l1_beta
|
|
|
|
def _get_ground_truth(self):
|
|
raise NotImplementedError()
|
|
|
|
def predict_proposals(self):
|
|
# pred_anchor_deltas: (L, N, ? Hi, Wi)
|
|
# anchors:(N, L, -1, B)
|
|
# here we loop over specific feature map, NOT images
|
|
proposals = []
|
|
anchors = self.anchors.transpose(0, 1)
|
|
for anchors_i, pred_anchor_deltas_i in zip(anchors, self.pred_anchor_deltas):
|
|
B = anchors_i.size(-1)
|
|
N, _, Hi, Wi = pred_anchor_deltas_i.shape
|
|
anchors_i = anchors_i.flatten(start_dim=0, end_dim=1)
|
|
pred_anchor_deltas_i = pred_anchor_deltas_i.view(N, -1, B, Hi, Wi).permute(0, 3, 4, 1, 2).reshape(-1, B)
|
|
proposals_i = self.box2box_transform.apply_deltas(pred_anchor_deltas_i, anchors_i)
|
|
# Append feature map proposals with shape (N, Hi*Wi*A, B)
|
|
proposals.append(proposals_i.view(N, -1, B))
|
|
proposals = torch.stack(proposals)
|
|
return proposals
|
|
|
|
def predict_objectness_logits(self):
|
|
"""
|
|
Returns:
|
|
pred_objectness_logits (list[Tensor]) -> (N, Hi*Wi*A).
|
|
"""
|
|
pred_objectness_logits = [
|
|
# Reshape: (N, A, Hi, Wi) -> (N, Hi, Wi, A) -> (N, Hi*Wi*A)
|
|
score.permute(0, 2, 3, 1).reshape(self.num_images, -1)
|
|
for score in self.pred_objectness_logits
|
|
]
|
|
return pred_objectness_logits
|
|
|
|
|
|
# Main Classes
|
|
class Conv2d(nn.Conv2d):
|
|
def __init__(self, *args, **kwargs):
|
|
norm = kwargs.pop("norm", None)
|
|
activation = kwargs.pop("activation", None)
|
|
super().__init__(*args, **kwargs)
|
|
|
|
self.norm = norm
|
|
self.activation = activation
|
|
|
|
def forward(self, x):
|
|
if x.numel() == 0 and self.training:
|
|
assert not isinstance(self.norm, nn.SyncBatchNorm)
|
|
if x.numel() == 0:
|
|
assert not isinstance(self.norm, nn.GroupNorm)
|
|
output_shape = [
|
|
(i + 2 * p - (di * (k - 1) + 1)) // s + 1
|
|
for i, p, di, k, s in zip(
|
|
x.shape[-2:],
|
|
self.padding,
|
|
self.dilation,
|
|
self.kernel_size,
|
|
self.stride,
|
|
)
|
|
]
|
|
output_shape = [x.shape[0], self.weight.shape[0]] + output_shape
|
|
empty = _NewEmptyTensorOp.apply(x, output_shape)
|
|
if self.training:
|
|
_dummy = sum(x.view(-1)[0] for x in self.parameters()) * 0.0
|
|
return empty + _dummy
|
|
else:
|
|
return empty
|
|
|
|
x = super().forward(x)
|
|
if self.norm is not None:
|
|
x = self.norm(x)
|
|
if self.activation is not None:
|
|
x = self.activation(x)
|
|
return x
|
|
|
|
|
|
class LastLevelMaxPool(nn.Module):
|
|
"""
|
|
This module is used in the original FPN to generate a downsampled P6 feature from P5.
|
|
"""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.num_levels = 1
|
|
self.in_feature = "p5"
|
|
|
|
def forward(self, x):
|
|
return [nn.functional.max_pool2d(x, kernel_size=1, stride=2, padding=0)]
|
|
|
|
|
|
class LastLevelP6P7(nn.Module):
|
|
"""
|
|
This module is used in RetinaNet to generate extra layers, P6 and P7 from C5 feature.
|
|
"""
|
|
|
|
def __init__(self, in_channels, out_channels):
|
|
super().__init__()
|
|
self.num_levels = 2
|
|
self.in_feature = "res5"
|
|
self.p6 = nn.Conv2d(in_channels, out_channels, 3, 2, 1)
|
|
self.p7 = nn.Conv2d(out_channels, out_channels, 3, 2, 1)
|
|
|
|
def forward(self, c5):
|
|
p6 = self.p6(c5)
|
|
p7 = self.p7(nn.functional.relu(p6))
|
|
return [p6, p7]
|
|
|
|
|
|
class BasicStem(nn.Module):
|
|
def __init__(self, in_channels=3, out_channels=64, norm="BN", caffe_maxpool=False):
|
|
super().__init__()
|
|
self.conv1 = Conv2d(
|
|
in_channels,
|
|
out_channels,
|
|
kernel_size=7,
|
|
stride=2,
|
|
padding=3,
|
|
bias=False,
|
|
norm=get_norm(norm, out_channels),
|
|
)
|
|
self.caffe_maxpool = caffe_maxpool
|
|
# use pad 1 instead of pad zero
|
|
|
|
def forward(self, x):
|
|
x = self.conv1(x)
|
|
x = nn.functional.relu_(x)
|
|
if self.caffe_maxpool:
|
|
x = nn.functional.max_pool2d(x, kernel_size=3, stride=2, padding=0, ceil_mode=True)
|
|
else:
|
|
x = nn.functional.max_pool2d(x, kernel_size=3, stride=2, padding=1)
|
|
return x
|
|
|
|
@property
|
|
def out_channels(self):
|
|
return self.conv1.out_channels
|
|
|
|
@property
|
|
def stride(self):
|
|
return 4 # = stride 2 conv -> stride 2 max pool
|
|
|
|
|
|
class ResNetBlockBase(nn.Module):
|
|
def __init__(self, in_channels, out_channels, stride):
|
|
super().__init__()
|
|
self.in_channels = in_channels
|
|
self.out_channels = out_channels
|
|
self.stride = stride
|
|
|
|
def freeze(self):
|
|
for p in self.parameters():
|
|
p.requires_grad = False
|
|
return self
|
|
|
|
|
|
class BottleneckBlock(ResNetBlockBase):
|
|
def __init__(
|
|
self,
|
|
in_channels,
|
|
out_channels,
|
|
bottleneck_channels,
|
|
stride=1,
|
|
num_groups=1,
|
|
norm="BN",
|
|
stride_in_1x1=False,
|
|
dilation=1,
|
|
):
|
|
super().__init__(in_channels, out_channels, stride)
|
|
|
|
if in_channels != out_channels:
|
|
self.shortcut = Conv2d(
|
|
in_channels,
|
|
out_channels,
|
|
kernel_size=1,
|
|
stride=stride,
|
|
bias=False,
|
|
norm=get_norm(norm, out_channels),
|
|
)
|
|
else:
|
|
self.shortcut = None
|
|
|
|
# The original MSRA ResNet models have stride in the first 1x1 conv
|
|
# The subsequent fb.torch.resnet and Caffe2 ResNe[X]t implementations have
|
|
# stride in the 3x3 conv
|
|
stride_1x1, stride_3x3 = (stride, 1) if stride_in_1x1 else (1, stride)
|
|
|
|
self.conv1 = Conv2d(
|
|
in_channels,
|
|
bottleneck_channels,
|
|
kernel_size=1,
|
|
stride=stride_1x1,
|
|
bias=False,
|
|
norm=get_norm(norm, bottleneck_channels),
|
|
)
|
|
|
|
self.conv2 = Conv2d(
|
|
bottleneck_channels,
|
|
bottleneck_channels,
|
|
kernel_size=3,
|
|
stride=stride_3x3,
|
|
padding=1 * dilation,
|
|
bias=False,
|
|
groups=num_groups,
|
|
dilation=dilation,
|
|
norm=get_norm(norm, bottleneck_channels),
|
|
)
|
|
|
|
self.conv3 = Conv2d(
|
|
bottleneck_channels,
|
|
out_channels,
|
|
kernel_size=1,
|
|
bias=False,
|
|
norm=get_norm(norm, out_channels),
|
|
)
|
|
|
|
def forward(self, x):
|
|
out = self.conv1(x)
|
|
out = nn.functional.relu_(out)
|
|
|
|
out = self.conv2(out)
|
|
out = nn.functional.relu_(out)
|
|
|
|
out = self.conv3(out)
|
|
|
|
if self.shortcut is not None:
|
|
shortcut = self.shortcut(x)
|
|
else:
|
|
shortcut = x
|
|
|
|
out += shortcut
|
|
out = nn.functional.relu_(out)
|
|
return out
|
|
|
|
|
|
class Backbone(nn.Module, metaclass=ABCMeta):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
@abstractmethod
|
|
def forward(self):
|
|
pass
|
|
|
|
@property
|
|
def size_divisibility(self):
|
|
"""
|
|
Some backbones require the input height and width to be divisible by a specific integer. This is
|
|
typically true for encoder / decoder type networks with lateral connection (e.g., FPN) for which feature maps need to match
|
|
dimension in the "bottom up" and "top down" paths. Set to 0 if no specific input size divisibility is required.
|
|
"""
|
|
return 0
|
|
|
|
def output_shape(self):
|
|
return {
|
|
name: ShapeSpec(
|
|
channels=self._out_feature_channels[name],
|
|
stride=self._out_feature_strides[name],
|
|
)
|
|
for name in self._out_features
|
|
}
|
|
|
|
@property
|
|
def out_features(self):
|
|
"""deprecated"""
|
|
return self._out_features
|
|
|
|
@property
|
|
def out_feature_strides(self):
|
|
"""deprecated"""
|
|
return {f: self._out_feature_strides[f] for f in self._out_features}
|
|
|
|
@property
|
|
def out_feature_channels(self):
|
|
"""deprecated"""
|
|
return {f: self._out_feature_channels[f] for f in self._out_features}
|
|
|
|
|
|
class ResNet(Backbone):
|
|
def __init__(self, stem, stages, num_classes=None, out_features=None):
|
|
"""
|
|
Args:
|
|
stem (nn.Module): a stem module
|
|
stages (list[list[ResNetBlock]]): several (typically 4) stages, each contains multiple :class:`ResNetBlockBase`.
|
|
num_classes (None or int): if None, will not perform classification.
|
|
out_features (list[str]): name of the layers whose outputs should be returned in forward. Can be anything in:
|
|
"stem", "linear", or "res2" ... If None, will return the output of the last layer.
|
|
"""
|
|
super(ResNet, self).__init__()
|
|
self.stem = stem
|
|
self.num_classes = num_classes
|
|
|
|
current_stride = self.stem.stride
|
|
self._out_feature_strides = {"stem": current_stride}
|
|
self._out_feature_channels = {"stem": self.stem.out_channels}
|
|
|
|
self.stages_and_names = []
|
|
for i, blocks in enumerate(stages):
|
|
for block in blocks:
|
|
assert isinstance(block, ResNetBlockBase), block
|
|
curr_channels = block.out_channels
|
|
stage = nn.Sequential(*blocks)
|
|
name = "res" + str(i + 2)
|
|
self.add_module(name, stage)
|
|
self.stages_and_names.append((stage, name))
|
|
self._out_feature_strides[name] = current_stride = int(
|
|
current_stride * np.prod([k.stride for k in blocks])
|
|
)
|
|
self._out_feature_channels[name] = blocks[-1].out_channels
|
|
|
|
if num_classes is not None:
|
|
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
|
|
self.linear = nn.Linear(curr_channels, num_classes)
|
|
|
|
# Sec 5.1 in "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour":
|
|
# "The 1000-way fully-connected layer is initialized by
|
|
# drawing weights from a zero-mean Gaussian with std of 0.01."
|
|
nn.init.normal_(self.linear.weight, stddev=0.01)
|
|
name = "linear"
|
|
|
|
if out_features is None:
|
|
out_features = [name]
|
|
self._out_features = out_features
|
|
assert len(self._out_features)
|
|
children = [x[0] for x in self.named_children()]
|
|
for out_feature in self._out_features:
|
|
assert out_feature in children, "Available children: {}".format(", ".join(children))
|
|
|
|
def forward(self, x):
|
|
outputs = {}
|
|
x = self.stem(x)
|
|
if "stem" in self._out_features:
|
|
outputs["stem"] = x
|
|
for stage, name in self.stages_and_names:
|
|
x = stage(x)
|
|
if name in self._out_features:
|
|
outputs[name] = x
|
|
if self.num_classes is not None:
|
|
x = self.avgpool(x)
|
|
x = self.linear(x)
|
|
if "linear" in self._out_features:
|
|
outputs["linear"] = x
|
|
return outputs
|
|
|
|
def output_shape(self):
|
|
return {
|
|
name: ShapeSpec(
|
|
channels=self._out_feature_channels[name],
|
|
stride=self._out_feature_strides[name],
|
|
)
|
|
for name in self._out_features
|
|
}
|
|
|
|
@staticmethod
|
|
def make_stage(
|
|
block_class,
|
|
num_blocks,
|
|
first_stride=None,
|
|
*,
|
|
in_channels,
|
|
out_channels,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Usually, layers that produce the same feature map spatial size
|
|
are defined as one "stage".
|
|
Under such definition, stride_per_block[1:] should all be 1.
|
|
"""
|
|
if first_stride is not None:
|
|
assert "stride" not in kwargs and "stride_per_block" not in kwargs
|
|
kwargs["stride_per_block"] = [first_stride] + [1] * (num_blocks - 1)
|
|
blocks = []
|
|
for i in range(num_blocks):
|
|
curr_kwargs = {}
|
|
for k, v in kwargs.items():
|
|
if k.endswith("_per_block"):
|
|
assert (
|
|
len(v) == num_blocks
|
|
), f"Argument '{k}' of make_stage should have the same length as num_blocks={num_blocks}."
|
|
newk = k[: -len("_per_block")]
|
|
assert newk not in kwargs, f"Cannot call make_stage with both {k} and {newk}!"
|
|
curr_kwargs[newk] = v[i]
|
|
else:
|
|
curr_kwargs[k] = v
|
|
|
|
blocks.append(block_class(in_channels=in_channels, out_channels=out_channels, **curr_kwargs))
|
|
in_channels = out_channels
|
|
|
|
return blocks
|
|
|
|
|
|
class ROIPooler(nn.Module):
|
|
"""
|
|
Region of interest feature map pooler that supports pooling from one or more
|
|
feature maps.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
output_size,
|
|
scales,
|
|
sampling_ratio,
|
|
canonical_box_size=224,
|
|
canonical_level=4,
|
|
):
|
|
super().__init__()
|
|
# assumption that stride is a power of 2.
|
|
min_level = -math.log2(scales[0])
|
|
max_level = -math.log2(scales[-1])
|
|
|
|
# a bunch of testing
|
|
assert math.isclose(min_level, int(min_level)) and math.isclose(max_level, int(max_level))
|
|
assert len(scales) == max_level - min_level + 1, "not pyramid"
|
|
assert 0 < min_level and min_level <= max_level
|
|
if isinstance(output_size, int):
|
|
output_size = (output_size, output_size)
|
|
assert len(output_size) == 2 and isinstance(output_size[0], int) and isinstance(output_size[1], int)
|
|
if len(scales) > 1:
|
|
assert min_level <= canonical_level and canonical_level <= max_level
|
|
assert canonical_box_size > 0
|
|
|
|
self.output_size = output_size
|
|
self.min_level = int(min_level)
|
|
self.max_level = int(max_level)
|
|
self.level_poolers = nn.ModuleList(RoIPool(output_size, spatial_scale=scale) for scale in scales)
|
|
self.canonical_level = canonical_level
|
|
self.canonical_box_size = canonical_box_size
|
|
|
|
def forward(self, feature_maps, boxes):
|
|
"""
|
|
Args:
|
|
feature_maps: List[torch.Tensor(N,C,W,H)]
|
|
box_lists: list[torch.Tensor])
|
|
Returns:
|
|
A tensor of shape(N*B, Channels, output_size, output_size)
|
|
"""
|
|
x = list(feature_maps.values())
|
|
num_level_assignments = len(self.level_poolers)
|
|
assert len(x) == num_level_assignments and len(boxes) == x[0].size(0)
|
|
|
|
pooler_fmt_boxes = convert_boxes_to_pooler_format(boxes)
|
|
|
|
if num_level_assignments == 1:
|
|
return self.level_poolers[0](x[0], pooler_fmt_boxes)
|
|
|
|
level_assignments = assign_boxes_to_levels(
|
|
boxes,
|
|
self.min_level,
|
|
self.max_level,
|
|
self.canonical_box_size,
|
|
self.canonical_level,
|
|
)
|
|
|
|
num_boxes = len(pooler_fmt_boxes)
|
|
num_channels = x[0].shape[1]
|
|
output_size = self.output_size[0]
|
|
|
|
dtype, device = x[0].dtype, x[0].device
|
|
output = torch.zeros(
|
|
(num_boxes, num_channels, output_size, output_size),
|
|
dtype=dtype,
|
|
device=device,
|
|
)
|
|
|
|
for level, (x_level, pooler) in enumerate(zip(x, self.level_poolers)):
|
|
inds = torch.nonzero(level_assignments == level).squeeze(1)
|
|
pooler_fmt_boxes_level = pooler_fmt_boxes[inds]
|
|
output[inds] = pooler(x_level, pooler_fmt_boxes_level)
|
|
|
|
return output
|
|
|
|
|
|
class ROIOutputs(object):
|
|
def __init__(self, cfg, training=False):
|
|
self.smooth_l1_beta = cfg.ROI_BOX_HEAD.SMOOTH_L1_BETA
|
|
self.box2box_transform = Box2BoxTransform(weights=cfg.ROI_BOX_HEAD.BBOX_REG_WEIGHTS)
|
|
self.training = training
|
|
self.score_thresh = cfg.ROI_HEADS.SCORE_THRESH_TEST
|
|
self.min_detections = cfg.MIN_DETECTIONS
|
|
self.max_detections = cfg.MAX_DETECTIONS
|
|
|
|
nms_thresh = cfg.ROI_HEADS.NMS_THRESH_TEST
|
|
if not isinstance(nms_thresh, list):
|
|
nms_thresh = [nms_thresh]
|
|
self.nms_thresh = nms_thresh
|
|
|
|
def _predict_boxes(self, proposals, box_deltas, preds_per_image):
|
|
num_pred = box_deltas.size(0)
|
|
B = proposals[0].size(-1)
|
|
K = box_deltas.size(-1) // B
|
|
box_deltas = box_deltas.view(num_pred * K, B)
|
|
proposals = torch.cat(proposals, dim=0).unsqueeze(-2).expand(num_pred, K, B)
|
|
proposals = proposals.reshape(-1, B)
|
|
boxes = self.box2box_transform.apply_deltas(box_deltas, proposals)
|
|
return boxes.view(num_pred, K * B).split(preds_per_image, dim=0)
|
|
|
|
def _predict_objs(self, obj_logits, preds_per_image):
|
|
probs = nn.functional.softmax(obj_logits, dim=-1)
|
|
probs = probs.split(preds_per_image, dim=0)
|
|
return probs
|
|
|
|
def _predict_attrs(self, attr_logits, preds_per_image):
|
|
attr_logits = attr_logits[..., :-1].softmax(-1)
|
|
attr_probs, attrs = attr_logits.max(-1)
|
|
return attr_probs.split(preds_per_image, dim=0), attrs.split(preds_per_image, dim=0)
|
|
|
|
@torch.no_grad()
|
|
def inference(
|
|
self,
|
|
obj_logits,
|
|
attr_logits,
|
|
box_deltas,
|
|
pred_boxes,
|
|
features,
|
|
sizes,
|
|
scales=None,
|
|
):
|
|
# only the pred boxes is the
|
|
preds_per_image = [p.size(0) for p in pred_boxes]
|
|
boxes_all = self._predict_boxes(pred_boxes, box_deltas, preds_per_image)
|
|
obj_scores_all = self._predict_objs(obj_logits, preds_per_image) # list of length N
|
|
attr_probs_all, attrs_all = self._predict_attrs(attr_logits, preds_per_image)
|
|
features = features.split(preds_per_image, dim=0)
|
|
|
|
# fun for each image too, also I can experiment and do multiple images
|
|
final_results = []
|
|
zipped = zip(boxes_all, obj_scores_all, attr_probs_all, attrs_all, sizes)
|
|
for i, (boxes, obj_scores, attr_probs, attrs, size) in enumerate(zipped):
|
|
for nms_t in self.nms_thresh:
|
|
outputs = do_nms(
|
|
boxes,
|
|
obj_scores,
|
|
size,
|
|
self.score_thresh,
|
|
nms_t,
|
|
self.min_detections,
|
|
self.max_detections,
|
|
)
|
|
if outputs is not None:
|
|
max_boxes, max_scores, classes, ids = outputs
|
|
break
|
|
|
|
if scales is not None:
|
|
scale_yx = scales[i]
|
|
max_boxes[:, 0::2] *= scale_yx[1]
|
|
max_boxes[:, 1::2] *= scale_yx[0]
|
|
|
|
final_results.append(
|
|
(
|
|
max_boxes,
|
|
classes,
|
|
max_scores,
|
|
attrs[ids],
|
|
attr_probs[ids],
|
|
features[i][ids],
|
|
)
|
|
)
|
|
boxes, classes, class_probs, attrs, attr_probs, roi_features = map(list, zip(*final_results))
|
|
return boxes, classes, class_probs, attrs, attr_probs, roi_features
|
|
|
|
def training(self, obj_logits, attr_logits, box_deltas, pred_boxes, features, sizes):
|
|
pass
|
|
|
|
def __call__(
|
|
self,
|
|
obj_logits,
|
|
attr_logits,
|
|
box_deltas,
|
|
pred_boxes,
|
|
features,
|
|
sizes,
|
|
scales=None,
|
|
):
|
|
if self.training:
|
|
raise NotImplementedError()
|
|
return self.inference(
|
|
obj_logits,
|
|
attr_logits,
|
|
box_deltas,
|
|
pred_boxes,
|
|
features,
|
|
sizes,
|
|
scales=scales,
|
|
)
|
|
|
|
|
|
class Res5ROIHeads(nn.Module):
|
|
"""
|
|
ROIHeads perform all per-region computation in an R-CNN.
|
|
It contains logic of cropping the regions, extract per-region features
|
|
(by the res-5 block in this case), and make per-region predictions.
|
|
"""
|
|
|
|
def __init__(self, cfg, input_shape):
|
|
super().__init__()
|
|
self.batch_size_per_image = cfg.RPN.BATCH_SIZE_PER_IMAGE
|
|
self.positive_sample_fraction = cfg.ROI_HEADS.POSITIVE_FRACTION
|
|
self.in_features = cfg.ROI_HEADS.IN_FEATURES
|
|
self.num_classes = cfg.ROI_HEADS.NUM_CLASSES
|
|
self.proposal_append_gt = cfg.ROI_HEADS.PROPOSAL_APPEND_GT
|
|
self.feature_strides = {k: v.stride for k, v in input_shape.items()}
|
|
self.feature_channels = {k: v.channels for k, v in input_shape.items()}
|
|
self.cls_agnostic_bbox_reg = cfg.ROI_BOX_HEAD.CLS_AGNOSTIC_BBOX_REG
|
|
self.stage_channel_factor = 2**3 # res5 is 8x res2
|
|
self.out_channels = cfg.RESNETS.RES2_OUT_CHANNELS * self.stage_channel_factor
|
|
|
|
# self.proposal_matcher = Matcher(
|
|
# cfg.ROI_HEADS.IOU_THRESHOLDS,
|
|
# cfg.ROI_HEADS.IOU_LABELS,
|
|
# allow_low_quality_matches=False,
|
|
# )
|
|
|
|
pooler_resolution = cfg.ROI_BOX_HEAD.POOLER_RESOLUTION
|
|
pooler_scales = (1.0 / self.feature_strides[self.in_features[0]],)
|
|
sampling_ratio = cfg.ROI_BOX_HEAD.POOLER_SAMPLING_RATIO
|
|
res5_halve = cfg.ROI_BOX_HEAD.RES5HALVE
|
|
use_attr = cfg.ROI_BOX_HEAD.ATTR
|
|
num_attrs = cfg.ROI_BOX_HEAD.NUM_ATTRS
|
|
|
|
self.pooler = ROIPooler(
|
|
output_size=pooler_resolution,
|
|
scales=pooler_scales,
|
|
sampling_ratio=sampling_ratio,
|
|
)
|
|
|
|
self.res5 = self._build_res5_block(cfg)
|
|
if not res5_halve:
|
|
"""
|
|
Modifications for VG in RoI heads:
|
|
1. Change the stride of conv1 and shortcut in Res5.Block1 from 2 to 1
|
|
2. Modifying all conv2 with (padding: 1 --> 2) and (dilation: 1 --> 2)
|
|
"""
|
|
self.res5[0].conv1.stride = (1, 1)
|
|
self.res5[0].shortcut.stride = (1, 1)
|
|
for i in range(3):
|
|
self.res5[i].conv2.padding = (2, 2)
|
|
self.res5[i].conv2.dilation = (2, 2)
|
|
|
|
self.box_predictor = FastRCNNOutputLayers(
|
|
self.out_channels,
|
|
self.num_classes,
|
|
self.cls_agnostic_bbox_reg,
|
|
use_attr=use_attr,
|
|
num_attrs=num_attrs,
|
|
)
|
|
|
|
def _build_res5_block(self, cfg):
|
|
stage_channel_factor = self.stage_channel_factor # res5 is 8x res2
|
|
num_groups = cfg.RESNETS.NUM_GROUPS
|
|
width_per_group = cfg.RESNETS.WIDTH_PER_GROUP
|
|
bottleneck_channels = num_groups * width_per_group * stage_channel_factor
|
|
out_channels = self.out_channels
|
|
stride_in_1x1 = cfg.RESNETS.STRIDE_IN_1X1
|
|
norm = cfg.RESNETS.NORM
|
|
|
|
blocks = ResNet.make_stage(
|
|
BottleneckBlock,
|
|
3,
|
|
first_stride=2,
|
|
in_channels=out_channels // 2,
|
|
bottleneck_channels=bottleneck_channels,
|
|
out_channels=out_channels,
|
|
num_groups=num_groups,
|
|
norm=norm,
|
|
stride_in_1x1=stride_in_1x1,
|
|
)
|
|
return nn.Sequential(*blocks)
|
|
|
|
def _shared_roi_transform(self, features, boxes):
|
|
x = self.pooler(features, boxes)
|
|
return self.res5(x)
|
|
|
|
def forward(self, features, proposal_boxes, gt_boxes=None):
|
|
if self.training:
|
|
"""
|
|
see https://github.com/airsplay/py-bottom-up-attention/\
|
|
blob/master/detectron2/modeling/roi_heads/roi_heads.py
|
|
"""
|
|
raise NotImplementedError()
|
|
|
|
assert not proposal_boxes[0].requires_grad
|
|
box_features = self._shared_roi_transform(features, proposal_boxes)
|
|
feature_pooled = box_features.mean(dim=[2, 3]) # pooled to 1x1
|
|
obj_logits, attr_logits, pred_proposal_deltas = self.box_predictor(feature_pooled)
|
|
return obj_logits, attr_logits, pred_proposal_deltas, feature_pooled
|
|
|
|
|
|
class AnchorGenerator(nn.Module):
|
|
"""
|
|
For a set of image sizes and feature maps, computes a set of anchors.
|
|
"""
|
|
|
|
def __init__(self, cfg, input_shape: List[ShapeSpec]):
|
|
super().__init__()
|
|
sizes = cfg.ANCHOR_GENERATOR.SIZES
|
|
aspect_ratios = cfg.ANCHOR_GENERATOR.ASPECT_RATIOS
|
|
self.strides = [x.stride for x in input_shape]
|
|
self.offset = cfg.ANCHOR_GENERATOR.OFFSET
|
|
assert 0.0 <= self.offset < 1.0, self.offset
|
|
|
|
"""
|
|
sizes (list[list[int]]): sizes[i] is the list of anchor sizes for feat map i
|
|
1. given in absolute lengths in units of the input image;
|
|
2. they do not dynamically scale if the input image size changes.
|
|
aspect_ratios (list[list[float]])
|
|
strides (list[int]): stride of each input feature.
|
|
"""
|
|
|
|
self.num_features = len(self.strides)
|
|
self.cell_anchors = nn.ParameterList(self._calculate_anchors(sizes, aspect_ratios))
|
|
self._spacial_feat_dim = 4
|
|
|
|
def _calculate_anchors(self, sizes, aspect_ratios):
|
|
# If one size (or aspect ratio) is specified and there are multiple feature
|
|
# maps, then we "broadcast" anchors of that single size (or aspect ratio)
|
|
if len(sizes) == 1:
|
|
sizes *= self.num_features
|
|
if len(aspect_ratios) == 1:
|
|
aspect_ratios *= self.num_features
|
|
assert self.num_features == len(sizes)
|
|
assert self.num_features == len(aspect_ratios)
|
|
|
|
cell_anchors = [self.generate_cell_anchors(s, a).float() for s, a in zip(sizes, aspect_ratios)]
|
|
|
|
return cell_anchors
|
|
|
|
@property
|
|
def box_dim(self):
|
|
return self._spacial_feat_dim
|
|
|
|
@property
|
|
def num_cell_anchors(self):
|
|
"""
|
|
Returns:
|
|
list[int]: Each int is the number of anchors at every pixel location, on that feature map.
|
|
"""
|
|
return [len(cell_anchors) for cell_anchors in self.cell_anchors]
|
|
|
|
def grid_anchors(self, grid_sizes):
|
|
anchors = []
|
|
for size, stride, base_anchors in zip(grid_sizes, self.strides, self.cell_anchors):
|
|
shift_x, shift_y = _create_grid_offsets(size, stride, self.offset, base_anchors.device)
|
|
shifts = torch.stack((shift_x, shift_y, shift_x, shift_y), dim=1)
|
|
|
|
anchors.append((shifts.view(-1, 1, 4) + base_anchors.view(1, -1, 4)).reshape(-1, 4))
|
|
|
|
return anchors
|
|
|
|
def generate_cell_anchors(self, sizes=(32, 64, 128, 256, 512), aspect_ratios=(0.5, 1, 2)):
|
|
"""
|
|
anchors are continuous geometric rectangles
|
|
centered on one feature map point sample.
|
|
We can later build the set of anchors
|
|
for the entire feature map by tiling these tensors
|
|
"""
|
|
|
|
anchors = []
|
|
for size in sizes:
|
|
area = size**2.0
|
|
for aspect_ratio in aspect_ratios:
|
|
w = math.sqrt(area / aspect_ratio)
|
|
h = aspect_ratio * w
|
|
x0, y0, x1, y1 = -w / 2.0, -h / 2.0, w / 2.0, h / 2.0
|
|
anchors.append([x0, y0, x1, y1])
|
|
return nn.Parameter(torch.tensor(anchors))
|
|
|
|
def forward(self, features):
|
|
"""
|
|
Args:
|
|
features List[torch.Tensor]: list of feature maps on which to generate anchors.
|
|
Returns:
|
|
torch.Tensor: a list of #image elements.
|
|
"""
|
|
num_images = features[0].size(0)
|
|
grid_sizes = [feature_map.shape[-2:] for feature_map in features]
|
|
anchors_over_all_feature_maps = self.grid_anchors(grid_sizes)
|
|
anchors_over_all_feature_maps = torch.stack(anchors_over_all_feature_maps)
|
|
return anchors_over_all_feature_maps.unsqueeze(0).repeat_interleave(num_images, dim=0)
|
|
|
|
|
|
class RPNHead(nn.Module):
|
|
"""
|
|
RPN classification and regression heads. Uses a 3x3 conv to produce a shared
|
|
hidden state from which one 1x1 conv predicts objectness logits for each anchor
|
|
and a second 1x1 conv predicts bounding-box deltas specifying how to deform
|
|
each anchor into an object proposal.
|
|
"""
|
|
|
|
def __init__(self, cfg, input_shape: List[ShapeSpec]):
|
|
super().__init__()
|
|
|
|
# Standard RPN is shared across levels:
|
|
in_channels = [s.channels for s in input_shape]
|
|
assert len(set(in_channels)) == 1, "Each level must have the same channel!"
|
|
in_channels = in_channels[0]
|
|
|
|
anchor_generator = AnchorGenerator(cfg, input_shape)
|
|
num_cell_anchors = anchor_generator.num_cell_anchors
|
|
box_dim = anchor_generator.box_dim
|
|
assert len(set(num_cell_anchors)) == 1, "Each level must have the same number of cell anchors"
|
|
num_cell_anchors = num_cell_anchors[0]
|
|
|
|
if cfg.PROPOSAL_GENERATOR.HIDDEN_CHANNELS == -1:
|
|
hid_channels = in_channels
|
|
else:
|
|
hid_channels = cfg.PROPOSAL_GENERATOR.HIDDEN_CHANNELS
|
|
# Modifications for VG in RPN (modeling/proposal_generator/rpn.py)
|
|
# Use hidden dim instead fo the same dim as Res4 (in_channels)
|
|
|
|
# 3x3 conv for the hidden representation
|
|
self.conv = nn.Conv2d(in_channels, hid_channels, kernel_size=3, stride=1, padding=1)
|
|
# 1x1 conv for predicting objectness logits
|
|
self.objectness_logits = nn.Conv2d(hid_channels, num_cell_anchors, kernel_size=1, stride=1)
|
|
# 1x1 conv for predicting box2box transform deltas
|
|
self.anchor_deltas = nn.Conv2d(hid_channels, num_cell_anchors * box_dim, kernel_size=1, stride=1)
|
|
|
|
for layer in [self.conv, self.objectness_logits, self.anchor_deltas]:
|
|
nn.init.normal_(layer.weight, std=0.01)
|
|
nn.init.constant_(layer.bias, 0)
|
|
|
|
def forward(self, features):
|
|
"""
|
|
Args:
|
|
features (list[Tensor]): list of feature maps
|
|
"""
|
|
pred_objectness_logits = []
|
|
pred_anchor_deltas = []
|
|
for x in features:
|
|
t = nn.functional.relu(self.conv(x))
|
|
pred_objectness_logits.append(self.objectness_logits(t))
|
|
pred_anchor_deltas.append(self.anchor_deltas(t))
|
|
return pred_objectness_logits, pred_anchor_deltas
|
|
|
|
|
|
class RPN(nn.Module):
|
|
"""
|
|
Region Proposal Network, introduced by the Faster R-CNN paper.
|
|
"""
|
|
|
|
def __init__(self, cfg, input_shape: Dict[str, ShapeSpec]):
|
|
super().__init__()
|
|
|
|
self.min_box_side_len = cfg.PROPOSAL_GENERATOR.MIN_SIZE
|
|
self.in_features = cfg.RPN.IN_FEATURES
|
|
self.nms_thresh = cfg.RPN.NMS_THRESH
|
|
self.batch_size_per_image = cfg.RPN.BATCH_SIZE_PER_IMAGE
|
|
self.positive_fraction = cfg.RPN.POSITIVE_FRACTION
|
|
self.smooth_l1_beta = cfg.RPN.SMOOTH_L1_BETA
|
|
self.loss_weight = cfg.RPN.LOSS_WEIGHT
|
|
|
|
self.pre_nms_topk = {
|
|
True: cfg.RPN.PRE_NMS_TOPK_TRAIN,
|
|
False: cfg.RPN.PRE_NMS_TOPK_TEST,
|
|
}
|
|
self.post_nms_topk = {
|
|
True: cfg.RPN.POST_NMS_TOPK_TRAIN,
|
|
False: cfg.RPN.POST_NMS_TOPK_TEST,
|
|
}
|
|
self.boundary_threshold = cfg.RPN.BOUNDARY_THRESH
|
|
|
|
self.anchor_generator = AnchorGenerator(cfg, [input_shape[f] for f in self.in_features])
|
|
self.box2box_transform = Box2BoxTransform(weights=cfg.RPN.BBOX_REG_WEIGHTS)
|
|
self.anchor_matcher = Matcher(
|
|
cfg.RPN.IOU_THRESHOLDS,
|
|
cfg.RPN.IOU_LABELS,
|
|
allow_low_quality_matches=True,
|
|
)
|
|
self.rpn_head = RPNHead(cfg, [input_shape[f] for f in self.in_features])
|
|
|
|
def training(self, images, image_shapes, features, gt_boxes):
|
|
pass
|
|
|
|
def inference(self, outputs, images, image_shapes, features, gt_boxes=None):
|
|
outputs = find_top_rpn_proposals(
|
|
outputs.predict_proposals(),
|
|
outputs.predict_objectness_logits(),
|
|
images,
|
|
image_shapes,
|
|
self.nms_thresh,
|
|
self.pre_nms_topk[self.training],
|
|
self.post_nms_topk[self.training],
|
|
self.min_box_side_len,
|
|
self.training,
|
|
)
|
|
|
|
results = []
|
|
for img in outputs:
|
|
im_boxes, img_box_logits = img
|
|
img_box_logits, inds = img_box_logits.sort(descending=True)
|
|
im_boxes = im_boxes[inds]
|
|
results.append((im_boxes, img_box_logits))
|
|
|
|
(proposal_boxes, logits) = tuple(map(list, zip(*results)))
|
|
return proposal_boxes, logits
|
|
|
|
def forward(self, images, image_shapes, features, gt_boxes=None):
|
|
"""
|
|
Args:
|
|
images (torch.Tensor): input images of length `N`
|
|
features (dict[str: Tensor])
|
|
gt_instances
|
|
"""
|
|
# features is dict, key = block level, v = feature_map
|
|
features = [features[f] for f in self.in_features]
|
|
pred_objectness_logits, pred_anchor_deltas = self.rpn_head(features)
|
|
anchors = self.anchor_generator(features)
|
|
outputs = RPNOutputs(
|
|
self.box2box_transform,
|
|
self.anchor_matcher,
|
|
self.batch_size_per_image,
|
|
self.positive_fraction,
|
|
images,
|
|
pred_objectness_logits,
|
|
pred_anchor_deltas,
|
|
anchors,
|
|
self.boundary_threshold,
|
|
gt_boxes,
|
|
self.smooth_l1_beta,
|
|
)
|
|
# For RPN-only models, the proposals are the final output
|
|
|
|
if self.training:
|
|
raise NotImplementedError()
|
|
return self.training(outputs, images, image_shapes, features, gt_boxes)
|
|
else:
|
|
return self.inference(outputs, images, image_shapes, features, gt_boxes)
|
|
|
|
|
|
class FastRCNNOutputLayers(nn.Module):
|
|
"""
|
|
Two linear layers for predicting Fast R-CNN outputs:
|
|
(1) proposal-to-detection box regression deltas
|
|
(2) classification scores
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
input_size,
|
|
num_classes,
|
|
cls_agnostic_bbox_reg,
|
|
box_dim=4,
|
|
use_attr=False,
|
|
num_attrs=-1,
|
|
):
|
|
"""
|
|
Args:
|
|
input_size (int): channels, or (channels, height, width)
|
|
num_classes (int)
|
|
cls_agnostic_bbox_reg (bool)
|
|
box_dim (int)
|
|
"""
|
|
super().__init__()
|
|
|
|
if not isinstance(input_size, int):
|
|
input_size = np.prod(input_size)
|
|
|
|
# (do + 1 for background class)
|
|
self.cls_score = nn.Linear(input_size, num_classes + 1)
|
|
num_bbox_reg_classes = 1 if cls_agnostic_bbox_reg else num_classes
|
|
self.bbox_pred = nn.Linear(input_size, num_bbox_reg_classes * box_dim)
|
|
|
|
self.use_attr = use_attr
|
|
if use_attr:
|
|
"""
|
|
Modifications for VG in RoI heads
|
|
Embedding: {num_classes + 1} --> {input_size // 8}
|
|
Linear: {input_size + input_size // 8} --> {input_size // 4}
|
|
Linear: {input_size // 4} --> {num_attrs + 1}
|
|
"""
|
|
self.cls_embedding = nn.Embedding(num_classes + 1, input_size // 8)
|
|
self.fc_attr = nn.Linear(input_size + input_size // 8, input_size // 4)
|
|
self.attr_score = nn.Linear(input_size // 4, num_attrs + 1)
|
|
|
|
nn.init.normal_(self.cls_score.weight, std=0.01)
|
|
nn.init.normal_(self.bbox_pred.weight, std=0.001)
|
|
for item in [self.cls_score, self.bbox_pred]:
|
|
nn.init.constant_(item.bias, 0)
|
|
|
|
def forward(self, roi_features):
|
|
if roi_features.dim() > 2:
|
|
roi_features = torch.flatten(roi_features, start_dim=1)
|
|
scores = self.cls_score(roi_features)
|
|
proposal_deltas = self.bbox_pred(roi_features)
|
|
if self.use_attr:
|
|
_, max_class = scores.max(-1) # [b, c] --> [b]
|
|
cls_emb = self.cls_embedding(max_class) # [b] --> [b, 256]
|
|
roi_features = torch.cat([roi_features, cls_emb], -1) # [b, 2048] + [b, 256] --> [b, 2304]
|
|
roi_features = self.fc_attr(roi_features)
|
|
roi_features = nn.functional.relu(roi_features)
|
|
attr_scores = self.attr_score(roi_features)
|
|
return scores, attr_scores, proposal_deltas
|
|
else:
|
|
return scores, proposal_deltas
|
|
|
|
|
|
class GeneralizedRCNN(nn.Module):
|
|
def __init__(self, cfg):
|
|
super().__init__()
|
|
|
|
self.device = torch.device(cfg.MODEL.DEVICE)
|
|
self.backbone = build_backbone(cfg)
|
|
self.proposal_generator = RPN(cfg, self.backbone.output_shape())
|
|
self.roi_heads = Res5ROIHeads(cfg, self.backbone.output_shape())
|
|
self.roi_outputs = ROIOutputs(cfg)
|
|
self.to(self.device)
|
|
|
|
@classmethod
|
|
def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
|
|
config = kwargs.pop("config", None)
|
|
state_dict = kwargs.pop("state_dict", None)
|
|
cache_dir = kwargs.pop("cache_dir", None)
|
|
from_tf = kwargs.pop("from_tf", False)
|
|
force_download = kwargs.pop("force_download", False)
|
|
resume_download = kwargs.pop("resume_download", False)
|
|
proxies = kwargs.pop("proxies", None)
|
|
local_files_only = kwargs.pop("local_files_only", False)
|
|
use_cdn = kwargs.pop("use_cdn", True)
|
|
|
|
# Load config if we don't provide a configuration
|
|
if not isinstance(config, Config):
|
|
config_path = config if config is not None else pretrained_model_name_or_path
|
|
# try:
|
|
config = Config.from_pretrained(
|
|
config_path,
|
|
cache_dir=cache_dir,
|
|
force_download=force_download,
|
|
resume_download=resume_download,
|
|
proxies=proxies,
|
|
local_files_only=local_files_only,
|
|
)
|
|
|
|
# Load model
|
|
if pretrained_model_name_or_path is not None:
|
|
if os.path.isdir(pretrained_model_name_or_path):
|
|
if os.path.isfile(os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)):
|
|
# Load from a PyTorch checkpoint
|
|
archive_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)
|
|
else:
|
|
raise EnvironmentError(
|
|
"Error no file named {} found in directory {} ".format(
|
|
WEIGHTS_NAME,
|
|
pretrained_model_name_or_path,
|
|
)
|
|
)
|
|
elif os.path.isfile(pretrained_model_name_or_path) or is_remote_url(pretrained_model_name_or_path):
|
|
archive_file = pretrained_model_name_or_path
|
|
elif os.path.isfile(pretrained_model_name_or_path + ".index"):
|
|
assert from_tf, "We found a TensorFlow checkpoint at {}, please set from_tf to True to load from this checkpoint".format(
|
|
pretrained_model_name_or_path + ".index"
|
|
)
|
|
archive_file = pretrained_model_name_or_path + ".index"
|
|
else:
|
|
archive_file = hf_bucket_url(
|
|
pretrained_model_name_or_path,
|
|
filename=WEIGHTS_NAME,
|
|
use_cdn=use_cdn,
|
|
)
|
|
|
|
try:
|
|
# Load from URL or cache if already cached
|
|
resolved_archive_file = cached_path(
|
|
archive_file,
|
|
cache_dir=cache_dir,
|
|
force_download=force_download,
|
|
proxies=proxies,
|
|
resume_download=resume_download,
|
|
local_files_only=local_files_only,
|
|
)
|
|
if resolved_archive_file is None:
|
|
raise EnvironmentError
|
|
except EnvironmentError:
|
|
msg = f"Can't load weights for '{pretrained_model_name_or_path}'."
|
|
raise EnvironmentError(msg)
|
|
|
|
if resolved_archive_file == archive_file:
|
|
print("loading weights file {}".format(archive_file))
|
|
else:
|
|
print("loading weights file {} from cache at {}".format(archive_file, resolved_archive_file))
|
|
else:
|
|
resolved_archive_file = None
|
|
|
|
# Instantiate model.
|
|
model = cls(config)
|
|
|
|
if state_dict is None:
|
|
try:
|
|
try:
|
|
state_dict = torch.load(resolved_archive_file, map_location="cpu")
|
|
except Exception:
|
|
state_dict = load_checkpoint(resolved_archive_file)
|
|
|
|
except Exception:
|
|
raise OSError(
|
|
"Unable to load weights from pytorch checkpoint file. "
|
|
"If you tried to load a PyTorch model from a TF 2.0 checkpoint, please set from_tf=True. "
|
|
)
|
|
|
|
missing_keys = []
|
|
unexpected_keys = []
|
|
error_msgs = []
|
|
|
|
# Convert old format to new format if needed from a PyTorch state_dict
|
|
old_keys = []
|
|
new_keys = []
|
|
for key in state_dict.keys():
|
|
new_key = None
|
|
if "gamma" in key:
|
|
new_key = key.replace("gamma", "weight")
|
|
if "beta" in key:
|
|
new_key = key.replace("beta", "bias")
|
|
if new_key:
|
|
old_keys.append(key)
|
|
new_keys.append(new_key)
|
|
for old_key, new_key in zip(old_keys, new_keys):
|
|
state_dict[new_key] = state_dict.pop(old_key)
|
|
|
|
# copy state_dict so _load_from_state_dict can modify it
|
|
metadata = getattr(state_dict, "_metadata", None)
|
|
state_dict = state_dict.copy()
|
|
if metadata is not None:
|
|
state_dict._metadata = metadata
|
|
|
|
model_to_load = model
|
|
model_to_load.load_state_dict(state_dict)
|
|
|
|
if model.__class__.__name__ != model_to_load.__class__.__name__:
|
|
base_model_state_dict = model_to_load.state_dict().keys()
|
|
head_model_state_dict_without_base_prefix = [
|
|
key.split(cls.base_model_prefix + ".")[-1] for key in model.state_dict().keys()
|
|
]
|
|
missing_keys.extend(head_model_state_dict_without_base_prefix - base_model_state_dict)
|
|
|
|
if len(unexpected_keys) > 0:
|
|
print(
|
|
f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when"
|
|
f" initializing {model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are"
|
|
f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task or"
|
|
" with another architecture (e.g. initializing a BertForSequenceClassification model from a"
|
|
" BertForPreTraining model).\n- This IS NOT expected if you are initializing"
|
|
f" {model.__class__.__name__} from the checkpoint of a model that you expect to be exactly identical"
|
|
" (initializing a BertForSequenceClassification model from a BertForSequenceClassification model)."
|
|
)
|
|
else:
|
|
print(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n")
|
|
if len(missing_keys) > 0:
|
|
print(
|
|
f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at"
|
|
f" {pretrained_model_name_or_path} and are newly initialized: {missing_keys}\nYou should probably"
|
|
" TRAIN this model on a down-stream task to be able to use it for predictions and inference."
|
|
)
|
|
else:
|
|
print(
|
|
f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at"
|
|
f" {pretrained_model_name_or_path}.\nIf your task is similar to the task the model of the checkpoint"
|
|
f" was trained on, you can already use {model.__class__.__name__} for predictions without further"
|
|
" training."
|
|
)
|
|
if len(error_msgs) > 0:
|
|
raise RuntimeError(
|
|
"Error(s) in loading state_dict for {}:\n\t{}".format(
|
|
model.__class__.__name__, "\n\t".join(error_msgs)
|
|
)
|
|
)
|
|
# Set model in evaluation mode to deactivate DropOut modules by default
|
|
model.eval()
|
|
|
|
return model
|
|
|
|
def forward(
|
|
self,
|
|
images,
|
|
image_shapes,
|
|
gt_boxes=None,
|
|
proposals=None,
|
|
scales_yx=None,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
kwargs:
|
|
max_detections (int), return_tensors {"np", "pt", None}, padding {None,
|
|
"max_detections"}, pad_value (int), location = {"cuda", "cpu"}
|
|
"""
|
|
if self.training:
|
|
raise NotImplementedError()
|
|
return self.inference(
|
|
images=images,
|
|
image_shapes=image_shapes,
|
|
gt_boxes=gt_boxes,
|
|
proposals=proposals,
|
|
scales_yx=scales_yx,
|
|
**kwargs,
|
|
)
|
|
|
|
@torch.no_grad()
|
|
def inference(
|
|
self,
|
|
images,
|
|
image_shapes,
|
|
gt_boxes=None,
|
|
proposals=None,
|
|
scales_yx=None,
|
|
**kwargs,
|
|
):
|
|
# run images through backbone
|
|
original_sizes = image_shapes * scales_yx
|
|
features = self.backbone(images)
|
|
|
|
# generate proposals if none are available
|
|
if proposals is None:
|
|
proposal_boxes, _ = self.proposal_generator(images, image_shapes, features, gt_boxes)
|
|
else:
|
|
assert proposals is not None
|
|
|
|
# pool object features from either gt_boxes, or from proposals
|
|
obj_logits, attr_logits, box_deltas, feature_pooled = self.roi_heads(features, proposal_boxes, gt_boxes)
|
|
|
|
# prepare FRCNN Outputs and select top proposals
|
|
boxes, classes, class_probs, attrs, attr_probs, roi_features = self.roi_outputs(
|
|
obj_logits=obj_logits,
|
|
attr_logits=attr_logits,
|
|
box_deltas=box_deltas,
|
|
pred_boxes=proposal_boxes,
|
|
features=feature_pooled,
|
|
sizes=image_shapes,
|
|
scales=scales_yx,
|
|
)
|
|
|
|
# will we pad???
|
|
subset_kwargs = {
|
|
"max_detections": kwargs.get("max_detections", None),
|
|
"return_tensors": kwargs.get("return_tensors", None),
|
|
"pad_value": kwargs.get("pad_value", 0),
|
|
"padding": kwargs.get("padding", None),
|
|
}
|
|
preds_per_image = torch.tensor([p.size(0) for p in boxes])
|
|
boxes = pad_list_tensors(boxes, preds_per_image, **subset_kwargs)
|
|
classes = pad_list_tensors(classes, preds_per_image, **subset_kwargs)
|
|
class_probs = pad_list_tensors(class_probs, preds_per_image, **subset_kwargs)
|
|
attrs = pad_list_tensors(attrs, preds_per_image, **subset_kwargs)
|
|
attr_probs = pad_list_tensors(attr_probs, preds_per_image, **subset_kwargs)
|
|
roi_features = pad_list_tensors(roi_features, preds_per_image, **subset_kwargs)
|
|
subset_kwargs["padding"] = None
|
|
preds_per_image = pad_list_tensors(preds_per_image, None, **subset_kwargs)
|
|
sizes = pad_list_tensors(image_shapes, None, **subset_kwargs)
|
|
normalized_boxes = norm_box(boxes, original_sizes)
|
|
return OrderedDict(
|
|
{
|
|
"obj_ids": classes,
|
|
"obj_probs": class_probs,
|
|
"attr_ids": attrs,
|
|
"attr_probs": attr_probs,
|
|
"boxes": boxes,
|
|
"sizes": sizes,
|
|
"preds_per_image": preds_per_image,
|
|
"roi_features": roi_features,
|
|
"normalized_boxes": normalized_boxes,
|
|
}
|
|
)
|