def pre_init_with_sift(trainPhotoTourDataset): patch_size = 65 ON_GPU = True SIFT = SIFTNet(patch_size=patch_size, do_cuda=ON_GPU) SIFT.eval() if ON_GPU: SIFT.cuda() transformed = [] for img in trainPhotoTourDataset.data: transformed.append( np.expand_dims(cv2.resize(img.cpu().numpy(), (65, 65)), axis=0)) phototour_loader = data_utils.DataLoader( PhototourTrainingData(transformed), batch_size=256, shuffle=False) descriptors = [] pbar = tqdm(enumerate(phototour_loader)) for batch_idx, data_a in pbar: if ON_GPU: torch_patches = Variable(data_a.type(torch.FloatTensor).cuda(), volatile=True) else: torch_patches = Variable(data_a.type(torch.FloatTensor), volatile=True) res = SIFT(torch_patches) sift = np.round(512. * res.data.cpu().numpy()).astype(np.float32) descriptors.extend(sift) return np.array(descriptors)
patches = [] current_patch_idx = start_patch_idx for y in range(0, height, patch_h): patch_idxs.append([]) for x in range(0, width, patch_w): patch = np.array(img.crop( (x, y, x + patch_w, y + patch_h))).astype(np.uint8) patches.append(patch) patch_idxs[-1].append(current_patch_idx) current_patch_idx += 1 return np2torch(np.array(patches)), patch_idxs, patch_idxs[-1][-1] model = SIFTNet(patch_size=41) model.cuda() model.eval() try: input_img_fname = sys.argv[1] output_fname = sys.argv[2] except: print "Wrong input format. Try ./rank_scale_patches.py imgs/ref.png out.txt" sys.exit(1) patches, idxs, max_idx = read_patch_file(input_img_fname, patch_w=73, patch_h=41) bw_patches = patches.float().mean(dim=1, keepdim=True) print bw_patches.shape n_patches = bw_patches.size(0)