text: import torch
  from torch.utils.data import Dataset
  import os
  import numpy as np
  import cv2
  import skimage.io as io
  import torchvision.transforms as transforms
  import config
  import math
  from tqdm import tqdm
  from CAPS.caps_model_v2 import CAPSModel
  
  
  class HPatchDataset(Dataset):
      def __init__(self, imdir):
          self.transform = transforms.Compose([transforms.ToTensor(),
                                               transforms.Normalize(mean=(0.485, 0.456, 0.406),
                                                                    std=(0.229, 0.224, 0.225)),
                                               ])
          self.imfs = []
          for f in os.listdir(imdir):
              scene_dir = os.path.join(imdir, f)
              self.imfs.extend([os.path.join(scene_dir, '{}.ppm').format(ind) for ind in range(1, 7)])
  
      def __getitem__(self, item):
          imf = self.imfs[item]
          im = io.imread(imf)
          im_tensor = self.transform(im)
          c, h, w = im_tensor.shape
          div_num = 16
          if (h % div_num) != 0 or (w % div_num) != 0:
              nh = math.ceil(h / div_num) * div_num
              nw = math.ceil(w / div_num) * div_num
              bot_pad = nh - h
              rig_pad = nw - w
              m = torch.nn.ReflectionPad2d((0, rig_pad, 0, bot_pad))
              im_tensor = m(im_tensor)
              # print("padding h :{}, padding w:{}".format(bot_pad, rig_pad))
              # print("padding shape: {}".format(im_tensor.shape))
  
          # using sift keypoints
          sift = cv2.SIFT_create()
          gray = cv2.cvtColor(im, cv2.COLOR_RGB2GRAY)
          kpts = sift.detect(gray)
          kpts = np.array([[kp.pt[0], kp.pt[1]] for kp in kpts])
          coord = torch.from_numpy(kpts).float()
          out = {'im': im_tensor, 'coord': coord, 'imf': imf}
          return out
  
      def __len__(self):
          return len(self.imfs)
  
  
  if __name__ == '__main__':
      # example code for extracting features for HPatches dataset, SIFT keypoint is used
      args = config.get_args()
      device = 'cuda' if torch.cuda.is_available() else 'cpu'
  
      dataset = HPatchDataset(args.extract_img_dir)
      data_loader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=args.workers)
  
      model = CAPSModel(args)
  
      outdir = args.extract_out_dir
      os.makedirs(outdir, exist_ok=True)
  
      with torch.no_grad():
          for data in tqdm(data_loader):
              im = data['im'].to(device)
              img_path = data['imf'][0]
              coord = data['coord'].to(device)
              feat_f = model.extract_features(im, coord)
              desc = feat_f.squeeze(0).detach().cpu().numpy()
              kpt = coord.cpu().numpy().squeeze(0)
  
              save_folder = os.path.join(outdir, os.path.basename(os.path.dirname(img_path)))
              os.makedirs(save_folder, exist_ok=True)
              save_file = os.path.join(save_folder, "{}.caps_effiv3".format(os.path.basename(img_path)))
  
              with open(save_file, 'wb') as output_file:
                  np.savez(
                      output_file,
                      keypoints=kpt,
                      scores=[],
                      descriptors=desc)
  
  