def proposal_layer(self, rpn_cls_prob, rpn_bbox_pred, rpn_trans_param,
im_info):
if self.is_train:
pre_nms_top_n = self.config['train_rpn_pre_nms_top_n']
post_nms_top_n = self.config['train_rpn_post_nms_top_n']
nms_thresh = self.config['train_rpn_nms_thresh']
else:
pre_nms_top_n = self.config['test_rpn_pre_nms_top_n']
post_nms_top_n = self.config['test_rpn_post_nms_top_n']
nms_thresh = self.config['test_rpn_nms_thresh']
# Get the scores and bounding boxes
scores = rpn_cls_prob[:, :, :, self.num_anchors:]
rpn_bbox_pred = rpn_bbox_pred.view((-1, 4))
scores = scores.contiguous().view(-1, 1)
rpn_trans_param = rpn_trans_param.view((-1, 6))
proposals = bbox_transform_inv(self.anchors, rpn_bbox_pred)
proposals = clip_boxes(proposals, im_info[:2])
# Pick the top region proposals
scores, order = scores.view(-1).sort(descending=True)
if pre_nms_top_n > 0:
order = order[:pre_nms_top_n]
scores = scores[:pre_nms_top_n].view(-1, 1)
proposals = proposals[order.data, :]
trans_param = rpn_trans_param[order.data, :]
# Non-maximal suppression
keep = nms(torch.cat((proposals, scores), 1).data, nms_thresh)
# Pick th top region proposals after NMS
if post_nms_top_n > 0:
keep = keep[:post_nms_top_n]
proposals = proposals[keep, :]
scores = scores[keep,]
trans_param = trans_param[keep, :]
# Only support single image as input
batch_inds = Variable(
proposals.data.new(proposals.size(0), 1).zero_())
blob = torch.cat((batch_inds, proposals), 1)
return blob, scores, trans_param
def test_gallery(net, dataset, use_cuda, output_dir, thresh=0.):
"""test gallery images"""
with open('config.yml', 'r') as f:
config = yaml.load(f)
num_images = len(dataset)
all_boxes = [0 for _ in range(num_images)]
all_features = [0 for _ in range(num_images)]
start = time.time()
for i in range(num_images):
im, im_info, orig_shape = dataset.next()
im = im.transpose([0, 3, 1, 2])
with torch.no_grad():
if use_cuda:
im = Variable(torch.from_numpy(im).cuda())
else:
im = Variable(torch.from_numpy(im))
scores, bbox_pred, rois, features = net.forward(im, None, im_info)
boxes = rois[:, 1:5] / im_info[2]
scores = np.reshape(scores, [scores.shape[0], -1])
bbox_pred = np.reshape(bbox_pred, [bbox_pred.shape[0], -1])
if config['test_bbox_reg']:
# Apply bounding-box regression deltas
box_deltas = bbox_pred
pred_boxes = bbox_transform_inv(
torch.from_numpy(boxes), torch.from_numpy(box_deltas)).numpy()
pred_boxes = clip_boxes(pred_boxes, orig_shape)
else:
# Simply repeat the boxes, once for each class
pred_boxes = np.tile(boxes, (1, scores.shape[1]))
boxes = pred_boxes
# skip j = 0, because it's the background class
j = 1
inds = np.where(scores[:, j] > thresh)[0]
cls_scores = scores[inds, j]
cls_boxes = boxes[inds, j * 4:(j + 1) * 4]
cls_dets = np.hstack(
(cls_boxes, cls_scores[:, np.newaxis])).astype(np.float32,
copy=False)
keep = nms(torch.from_numpy(cls_dets),
config['test_nms']).numpy() if cls_dets.size > 0 else []
cls_dets = cls_dets[keep, :]
all_boxes[i] = cls_dets
all_features[i] = features[inds][keep]
end = time.time()
print('im_detect: {:d}/{:d} {:.3f}s'.format(i + 1, num_images,
(end - start) / (i + 1)))
det_file = os.path.join(output_dir, 'gboxes.pkl')
with open(det_file, 'wb') as f:
pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)
feature_file = os.path.join(output_dir, 'gfeatures.pkl')
with open(feature_file, 'wb') as f:
pickle.dump(all_features, f, pickle.HIGHEST_PROTOCOL)
return all_boxes, all_features
def test_gallery(net, dataloader, output_dir, thresh=0.):
"""test gallery images"""
with open('config.yml', 'r') as f:
config = yaml.load(f)
num_images = len(dataloader.dataset)
all_boxes = []
all_features = []
end = time.time()
time_cost = AverageMeter()
net.eval()
for i, data in enumerate(dataloader):
with torch.no_grad():
im, (orig_shape, im_info) = data
im = im.to(device)
im_info = im_info.numpy().squeeze(0)
orig_shape = [x.item() for x in orig_shape]
scores, bbox_pred, rois, features = net.forward(im, None, im_info)
boxes = rois[:, 1:5] / im_info[2]
scores = np.reshape(scores, [scores.shape[0], -1])
bbox_pred = np.reshape(bbox_pred, [bbox_pred.shape[0], -1])
if config['test_bbox_reg']:
# Apply bounding-box regression deltas
box_deltas = bbox_pred
pred_boxes = bbox_transform_inv(
torch.from_numpy(boxes), torch.from_numpy(box_deltas)).numpy()
pred_boxes = clip_boxes(pred_boxes, orig_shape)
else:
# Simply repeat the boxes, once for each class
pred_boxes = np.tile(boxes, (1, scores.shape[1]))
boxes = pred_boxes
# skip j = 0, because it's the background class
j = 1
inds = np.where(scores[:, j] > thresh)[0]
cls_scores = scores[inds, j]
cls_boxes = boxes[inds, j * 4:(j + 1) * 4]
cls_dets = np.hstack(
(cls_boxes, cls_scores[:, np.newaxis])).astype(np.float32,
copy=False)
keep = nms(torch.from_numpy(cls_dets),
config['test_nms']).numpy() if cls_dets.size > 0 else []
cls_dets = cls_dets[keep, :]
all_boxes.append(cls_dets)
all_features.append(features[inds][keep])
time_cost.update(time.time() - end)
end = time.time()
print('im_detect: {:d}/{:d} {:.3f}s'.format(i + 1, num_images,
time_cost.avg))
det_file = os.path.join(output_dir, 'gboxes.pkl')
with open(det_file, 'wb') as f:
pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)
feature_file = os.path.join(output_dir, 'gfeatures.pkl')
with open(feature_file, 'wb') as f:
pickle.dump(all_features, f, pickle.HIGHEST_PROTOCOL)
return all_boxes, all_features
def demo_detection(net, im_dir, images, use_cuda, thresh=.75):
with open('config.yml', 'r') as f:
config = yaml.load(f)
for im_name in images:
im_path = os.path.join(im_dir, im_name)
im, im_scale, orig_shape = pre_process_image(im_path, copy=True)
im_info = np.array([im.shape[1], im.shape[2], im_scale],
dtype=np.float32)
im = im.transpose([0, 3, 1, 2])
if use_cuda:
im = Variable(torch.from_numpy(im).cuda(), volatile=True)
else:
im = Variable(torch.from_numpy(im), volatile=True)
scores, bbox_pred, rois, _ = net.forward(im, None, im_info)
boxes = rois[:, 1:5] / im_info[2]
scores = np.reshape(scores, [scores.shape[0], -1])
bbox_pred = np.reshape(bbox_pred, [bbox_pred.shape[0], -1])
if config['test_bbox_reg']:
# Apply bounding-box regression deltas
box_deltas = bbox_pred
pred_boxes = bbox_transform_inv(
torch.from_numpy(boxes), torch.from_numpy(box_deltas)).numpy()
pred_boxes = clip_boxes(pred_boxes, orig_shape)
else:
# Simply repeat the boxes, once for each class
pred_boxes = np.tile(boxes, (1, scores.shape[1]))
boxes = pred_boxes
# skip j = 0, because it's the background class
j = 1
inds = np.where(scores[:, j] > thresh)[0]
cls_scores = scores[inds, j]
cls_boxes = boxes[inds, j * 4:(j + 1) * 4]
cls_dets = np.hstack((cls_boxes, cls_scores[:, np.newaxis])).astype(
np.float32, copy=False)
keep = nms(torch.from_numpy(cls_dets),
config['test_nms']).numpy() if cls_dets.size > 0 else []
cls_dets = cls_dets[keep, :]
if cls_dets is None:
print('There are no detections in image {}'.format(im_name))
continue
fig, ax = plt.subplots(figsize=(16, 9))
ax.imshow(plt.imread(im_path))
plt.axis('off')
for box in cls_dets:
x1, y1, x2, y2, score = box
ax.add_patch(plt.Rectangle((x1, y1), x2 - x1, y2 - y1, fill=False,
edgecolor='#66D9EF', linewidth=3.5))
ax.add_patch(plt.Rectangle((x1, y1), x2 - x1, y2 - y1, fill=False,
edgecolor='white', linewidth=1))
ax.text(x1 + 5, y1 - 15, '{:.2f}'.format(score),
bbox=dict(facecolor='#66D9EF', linewidth=0), fontsize=20,
color='white')
plt.tight_layout()
plt.show()
plt.close(fig)
def demo_search(net, im_dir, images, use_cuda, thresh=.75):
with open('config.yml', 'r') as f:
config = yaml.load(f)
q_name = 's15166.jpg'
q_roi = [29, 5, 164, 439] # x1, y1, h, w
x1, y1, h, w = q_roi
q_path = os.path.join(im_dir, q_name)
q_im, q_scale, _ = pre_process_image(q_path)
q_roi = np.array(q_roi) * q_scale
q_info = np.array([q_im.shape[1], q_im.shape[2], q_scale],
dtype=np.float32)
q_im = q_im.transpose([0, 3, 1, 2])
q_roi = np.hstack(([[0]], q_roi.reshape(1, 4)))
if use_cuda:
q_im = Variable(torch.from_numpy(q_im).cuda(), volatile=True)
q_roi = Variable(torch.from_numpy(q_roi).float().cuda())
else:
q_im = Variable(torch.from_numpy(q_im), volatile=True)
q_roi = Variable(torch.from_numpy(q_roi).float())
q_feat = net.forward(q_im, q_roi, q_info, 'query')[0]
# Show query
fig, ax = plt.subplots(figsize=(16, 9))
ax.imshow(plt.imread(q_path))
plt.axis('off')
ax.add_patch(plt.Rectangle((x1, y1), h, w, fill=False, edgecolor='#F92672',
linewidth=3.5))
ax.add_patch(plt.Rectangle((x1, y1), h, w, fill=False, edgecolor='white',
linewidth=1))
ax.text(x1 + 5, y1 - 15, '{}'.format('Query'),
bbox=dict(facecolor='#F92672', linewidth=0), fontsize=20,
color='white')
plt.tight_layout()
fig.savefig(os.path.join(im_dir, 'query.jpg'))
plt.show()
plt.close(fig)
# Get gallery images
images.remove(q_name)
for im_name in images:
im_path = os.path.join(im_dir, im_name)
im, im_scale, orig_shape = pre_process_image(im_path, copy=True)
im_info = np.array([im.shape[1], im.shape[2], im_scale],
dtype=np.float32)
im = im.transpose([0, 3, 1, 2])
if use_cuda:
im = Variable(torch.from_numpy(im).cuda(), volatile=True)
else:
im = Variable(torch.from_numpy(im), volatile=True)
scores, bbox_pred, rois, features = net.forward(im, None, im_info)
boxes = rois[:, 1:5] / im_info[2]
scores = np.reshape(scores, [scores.shape[0], -1])
bbox_pred = np.reshape(bbox_pred, [bbox_pred.shape[0], -1])
if config['test_bbox_reg']:
# Apply bounding-box regression deltas
box_deltas = bbox_pred
pred_boxes = bbox_transform_inv(
torch.from_numpy(boxes), torch.from_numpy(box_deltas)).numpy()
pred_boxes = clip_boxes(pred_boxes, orig_shape)
else:
# Simply repeat the boxes, once for each class
pred_boxes = np.tile(boxes, (1, scores.shape[1]))
boxes = pred_boxes
# skip j = 0, because it's the background class
j = 1
inds = np.where(scores[:, j] > thresh)[0]
cls_scores = scores[inds, j]
cls_boxes = boxes[inds, j * 4:(j + 1) * 4]
cls_dets = np.hstack((cls_boxes, cls_scores[:, np.newaxis])).astype(
np.float32, copy=False)
keep = nms(torch.from_numpy(cls_dets),
config['test_nms']).numpy() if cls_dets.size > 0 else []
cls_dets = cls_dets[keep, :]
features = features[inds][keep]
if cls_dets is None:
print('There are no detections in image {}'.format(im_name))
continue
similarities = features.dot(q_feat)
fig, ax = plt.subplots(figsize=(16, 9))
ax.imshow(plt.imread(im_path))
plt.axis('off')
# Set different colors for different ids
similarities_list = similarities.tolist()
max_sim = max(similarities_list)
similarities_list.remove(max_sim)
colors = {value: '#66D9EF' for value in similarities_list}
colors[max_sim] = '#4CAF50'
for box, sim in zip(cls_dets, similarities):
x1, y1, x2, y2, _ = box
ax.add_patch(plt.Rectangle((x1, y1), x2 - x1, y2 - y1, fill=False,
edgecolor=colors[sim], linewidth=3.5))
ax.add_patch(plt.Rectangle((x1, y1), x2 - x1, y2 - y1, fill=False,
edgecolor='white', linewidth=1))
ax.text(x1 + 5, y1 - 15, '{:.2f}'.format(sim),
bbox=dict(facecolor=colors[sim], linewidth=0), fontsize=20,
color='white')
plt.tight_layout()
fig.savefig(os.path.join(im_dir, 'result_' + im_name))
plt.show()
plt.close(fig)
