parser = argparse.ArgumentParser(description='Train a Fast R-CNN network')

parser.add_argument('--net', dest='net', help='fsod, qkv', default='fsod', type=str)

args = parser.parse_args()

if net == 'fsod':

model = FSOD(['bg', 'fg'], 50, pretrained=False, num_way=2, num_shot=n_shot)

elif net == 'multi':

model = MultiheadAttentionRCNN(['f', 'g'], 1, pretrained=False, num_way=2, num_shot=n_shot)

else:

raise Exception('model undefined')

model.create_architecture()

print("load checkpoint %s" % (load_path))

checkpoint = torch.load(load_path)

model.load_state_dict(checkpoint['model'])

if 'pooling_mode' in checkpoint.keys():

cfg.POOLING_MODE = checkpoint['pooling_mode']

model.eval()

cfg.CUDA = True

model.cuda()

im_data = torch.FloatTensor(1)

im_info = torch.FloatTensor(1)

num_boxes = torch.LongTensor(1)

gt_boxes = torch.FloatTensor(1)

support_ims = torch.FloatTensor(1)

im_data = im_data.cuda()

im_info = im_info.cuda()

num_boxes = num_boxes.cuda()

gt_boxes = gt_boxes.cuda()

support_ims = support_ims.cuda()

im_data = Variable(im_data)

im_info = Variable(im_info)

num_boxes = Variable(num_boxes)

gt_boxes = Variable(gt_boxes)

support_ims = Variable(support_ims)

support_data_all = np.zeros((n_of_shot, 3, support_im_size, support_im_size), dtype=np.float32)

for i, im in enumerate(im_list):

im = im[:,:,::-1] # rgb -> bgr

target_size = np.min(im.shape[0:2]) # don't change the size

im, _ = prep_im_for_blob(im, cfg.PIXEL_MEANS, target_size, cfg.TRAIN.MAX_SIZE)

_h, _w = im.shape[0], im.shape[1]

if _h > _w:

resize_scale = float(support_im_size) / float(_h)

unfit_size = int(_w * resize_scale)

im = cv2.resize(im, (unfit_size, support_im_size), interpolation=cv2.INTER_LINEAR)

else:

resize_scale = float(support_im_size) / float(_w)

unfit_size = int(_h * resize_scale)

im = cv2.resize(im, (support_im_size, unfit_size), interpolation=cv2.INTER_LINEAR)

h, w = im.shape[0], im.shape[1]

support_data_all[i, :, :h, :w] = np.transpose(im, (2, 0, 1))

support_data = torch.from_numpy(support_data_all).unsqueeze(0)

target_size = cfg.TRAIN.SCALES[0]

im_data, im_scale = prep_im_for_blob(im_data, cfg.PIXEL_MEANS, target_size, cfg.TRAIN.MAX_SIZE)

im_data = torch.from_numpy(im_data)

im_info = np.array([[im_data.shape[0], im_data.shape[1], im_scale]], dtype=np.float32)

im_info = torch.from_numpy(im_info)

gt_boxes = torch.from_numpy(np.array([0]))

num_boxes = torch.from_numpy(np.array([0]))

query = im_data.permute(2, 0, 1).contiguous().unsqueeze(0)

# support

# support_root_dir = 'datasets/supports'

# class_dir = 'horse'

# n_shot = 2

# im_paths = list(Path(os.path.join(support_root_dir, class_dir)).glob('*.jpg'))

CWD = os.getcwd()

print(support_im_paths)

n_shot = len(support_im_paths)

random.seed(0)

im_path_list = random.sample(support_im_paths, k=n_shot)

im_list = []

#fig = plt.figure(num=None, figsize=(8, 8), dpi=50, facecolor='w', edgecolor='k')

for i, im_path in enumerate(im_path_list):

im = Image.open(im_path)

im_list.append(np.asarray(im))

support_data = support_im_preprocess(im_list, cfg, 320, n_shot)

# query

im = np.asarray(Image.open(query_path))

im2show = im.copy()

query_data, im_info, gt_boxes, num_boxes = query_im_preprocess(im, cfg)

# prepare data

data = [query_data, im_info, gt_boxes, num_boxes, support_data]

im_data, im_info, num_boxes, gt_boxes, support_ims = prepare_variable()

with torch.no_grad():

im_data.resize_(data[0].size()).copy_(data[0])

im_info.resize_(data[1].size()).copy_(data[1])

gt_boxes.resize_(data[2].size()).copy_(data[2])

num_boxes.resize_(data[3].size()).copy_(data[3])

support_ims.resize_(data[4].size()).copy_(data[4])

# model

cfg_from_list(['ANCHOR_SCALES', '[4, 8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]'])

model_dir = os.path.join(CWD, 'models')

load_path = os.path.join(model_dir,'faster_rcnn_{}_{}_{}.pth'.format(1, 11, 34467))

model = get_model('multi', load_path, n_shot)

start_time = time.time()

rois, cls_prob, bbox_pred, \

rpn_loss_cls, rpn_loss_box, \

RCNN_loss_cls, RCNN_loss_bbox, \

rois_label = model(im_data, im_info, gt_boxes, num_boxes, support_ims, gt_boxes)

scores = cls_prob.data

boxes = rois.data[:, :, 1:5]

box_deltas = bbox_pred.data

if cfg.TRAIN.BBOX_NORMALIZE_TARGETS_PRECOMPUTED:

# Optionally normalize targets by a precomputed mean and stdev

box_deltas = box_deltas.view(-1, 4) * torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_STDS).cuda() \

+ torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_MEANS).cuda()

box_deltas = box_deltas.view(1, -1, 4)

pred_boxes = bbox_transform_inv(boxes, box_deltas, 1)

pred_boxes = clip_boxes(pred_boxes, im_info.data, 1)

# re-scale boxes to the origin img scale

pred_boxes /= data[1][0][2].item()

scores = scores.squeeze()

pred_boxes = pred_boxes.squeeze()

thresh = 0.05

inds = torch.nonzero(scores[:,1]>thresh).view(-1)

cls_scores = scores[:,1][inds]

_, order = torch.sort(cls_scores, 0, True)

cls_boxes = pred_boxes[inds, :]

cls_dets = torch.cat((cls_boxes, cls_scores.unsqueeze(1)), 1)

cls_dets = cls_dets[order]

keep = nms(cls_boxes[order, :], cls_scores[order], cfg.TEST.NMS)

cls_dets = cls_dets[keep.view(-1).long()]

for i in range(cls_dets.shape[0]):

w = cls_dets[i, 2] - cls_dets[i, 0]

h = cls_dets[i, 3] - cls_dets[i, 1]

if w > 0.5 * im2show.shape[1] or h > 0.5 * im2show.shape[0]:

cls_dets[i, 4] = 0

end_time = time.time()

im2show = vis_detections(im2show, ' ', cls_dets.cpu().numpy(), 0.5)

output_path = os.path.join(output_path_folder, 'result'+ str(cnt_shot) +'.jpg')

cv2.imwrite(output_path, im2show[:, :, ::-1])

print(cls_dets)