67 lines
2.2 KiB
Python
67 lines
2.2 KiB
Python
import os
|
|
import errno
|
|
import numpy as np
|
|
import torch
|
|
|
|
from detectron2.structures import Instances
|
|
from models.bua.layers.nms import nms
|
|
|
|
def save_features(output_file, features, boxes=None):
|
|
if boxes is None:
|
|
res = features
|
|
np.save(output_file, res)
|
|
else:
|
|
np.savez(output_file, x=features, bbox=boxes)
|
|
|
|
def mkdir(path):
|
|
try:
|
|
os.makedirs(path)
|
|
except OSError as e:
|
|
if e.errno != errno.EEXIST:
|
|
raise
|
|
|
|
def extractor_postprocess(boxes, scores, features_pooled, input_per_image, extractor):
|
|
"""
|
|
Resize the output instances.
|
|
The input images are often resized when entering an object detector.
|
|
As a result, we often need the outputs of the detector in a different
|
|
resolution from its inputs.
|
|
|
|
This function will resize the raw outputs of an R-CNN detector
|
|
to produce outputs according to the desired output resolution.
|
|
|
|
Args:
|
|
results (Instances): the raw outputs from the detector.
|
|
`results.image_size` contains the input image resolution the detector sees.
|
|
This object might be modified in-place.
|
|
output_height, output_width: the desired output resolution.
|
|
|
|
Returns:
|
|
Instances: the resized output from the model, based on the output resolution
|
|
"""
|
|
MIN_BOXES = extractor.MIN_BOXES
|
|
MAX_BOXES = extractor.MAX_BOXES
|
|
CONF_THRESH = extractor.CONF_THRESH
|
|
|
|
cur_device = scores.device
|
|
|
|
dets = boxes / input_per_image["im_scale"]
|
|
|
|
max_conf = torch.zeros((scores.shape[0])).to(cur_device)
|
|
|
|
for cls_ind in range(1, scores.shape[1]):
|
|
cls_scores = scores[:, cls_ind]
|
|
keep = nms(dets, cls_scores, 0.3)
|
|
max_conf[keep] = torch.where(cls_scores[keep] > max_conf[keep],
|
|
cls_scores[keep],
|
|
max_conf[keep])
|
|
|
|
keep_boxes = torch.nonzero(max_conf >= CONF_THRESH).flatten()
|
|
if len(keep_boxes) < MIN_BOXES:
|
|
keep_boxes = torch.argsort(max_conf, descending=True)[:MIN_BOXES]
|
|
elif len(keep_boxes) > MAX_BOXES:
|
|
keep_boxes = torch.argsort(max_conf, descending=True)[:MAX_BOXES]
|
|
image_feat = features_pooled[keep_boxes]
|
|
image_bboxes = dets[keep_boxes]
|
|
|
|
return image_feat, image_bboxes |