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)