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)