code: from cgi import test
  import torch
  from torchvision import transforms
  import torchvision.models as models
  from torch.utils import data
  from sklearn.neighbors import KNeighborsClassifier
  from sklearn.metrics.pairwise import cosine_similarity
  from scipy.spatial.distance import cdist
  
  from PIL import Image
  import pandas as pd
  from tqdm import tqdm
  from pathlib import Path
  import numpy as np
  import h5py
  import os
  import time
  
  class Identity(torch.nn.Module):
      def __init__(self):
          super(Identity, self).__init__()
          
      def forward(self, x):
          return x
  
  
  def get_valid_indices(train_path, test_path):
      train_df = pd.read_csv(train_path, index_col=False)
      test_study_ids = pd.read_csv(test_path, index_col=False)['study_id']
      train_df['study_id'] = train_df['filename'].astype(str).str[1:-4].astype(int)
      train_study_ids = train_df['study_id']
      contained = train_study_ids.isin(test_study_ids)
      not_present_indices = list(contained[contained == False].index)
      return not_present_indices
  
  class MIMICEncodingsDataset(data.Dataset):
      def __init__(self, encodings_path):
          encodings_file = h5py.File(encodings_path)
          self.encodings = encodings_file.get('encodings')
          self.reports = encodings_file.get('reports')
          self.indices = get_valid_indices('/deep/group/data/med-data/mimic_cxr_impressions.csv', '/deep/group/data/med-data/mimic-cxr-jpg-split/bootstrap_test/reports.csv')
  
      def __len__(self):
          return np.shape(self.encodings)[0]
          # return len(self.indices)
  
      def __getitem__(self, idx):
          return self.encodings[idx], self.reports[idx]
          # train_index = self.indices[idx]
          # return self.encodings[train_index], self.reports[train_index]
  
  
  class MIMICDataset(data.Dataset):
      """Represents an abstract HDF5 dataset.
      """
      def __init__(self, img_path, txt_path, size=None, transform=None):
          super().__init__()
          self.img_dset = h5py.File(img_path, 'r')['cxr_unprocessed']
          self.txt_dset = pd.read_csv(txt_path)['report']
          self.transform = transform
              
      def __len__(self):
          return len(self.txt_dset)
      
      def __getitem__(self, idx):
          if torch.is_tensor(idx):
              idx = idx.tolist()
              
          img = self.img_dset[idx] # np array, (320, 320)
          img = np.expand_dims(img, axis=0)
          img = np.repeat(img, 3, axis=0)
          txt = self.txt_dset[idx] # python str
  
          img = torch.from_numpy(img) # torch, (3, 320, 320)
          if self.transform:
              img = self.transform(img)
          sample = {'img': img, 'txt': txt }
          
          return sample
  
  class MIMICTestDataset(torch.utils.data.Dataset):
      def __init__(self, img_path, transform=None):
          super().__init__()
          self.img_dset = h5py.File(img_path, 'r')['cxr']
          self.transform = transform
  
      def __len__(self):
          return len(self.img_dset)
  
      def __getitem__(self, idx):
          img = self.img_dset[idx] # np array, (320, 320)
          img = np.expand_dims(img, axis=0)
          img = np.repeat(img, 3, axis=0)
  
          img = torch.from_numpy(img) # torch, (3, 320, 320)
          if self.transform:
              img = self.transform(img)
          return img
  
  class CheXpertDataset(data.Dataset):
      def __init__(self, img_path, transform=None):
          super().__init__()
  
          imgs_df = pd.read_csv(img_path)
          root_path = "/deep/group/CheXpert/CodaLab/"
          self.scale = 320
          self.transform = transform
          self.paths = []
          for _path in imgs_df["Path"]:
              if "view1" not in _path: continue # TODO: check how to aggregate studies, consider them independent??
              _pth = _path.replace("CheXpert-v1.0", "")
              _pth = Path(root_path+_pth)
              self.paths.append(_pth)
      def __len__(self):
          return len(self.paths)
      
      def __getitem__(self, idx):
          # img = Image.open(self.paths[idx]).resize((self.scale, self.scale)).convert('RGB')
          _np_img = np.asarray(Image.open(self.paths[idx]).resize((self.scale, self.scale), resample=Image.BICUBIC)) # these images all have diff sizes!
          _np_img = _np_img.astype('float32')
          img = np.expand_dims(_np_img, axis=0)
          img = np.repeat(img, 3, axis=0)
          img = torch.from_numpy(img)
  
          if self.transform:
              img = self.transform(img) # goes from H x W x C to C x H x W
          return img
  
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  
  # SET VARIABLES
  save_train_encodings = False
  retrieve_most_similar_train = True
  batch_size = 256
  chexpert_filepath = "/deep/group/CheXpert/CodaLab/test_image_paths.csv"
  cxr_filepath = '/deep/group/data/med-data/cxr.h5'
  txt_filepath = '/deep/group/data/med-data/mimic_cxr_impressions.csv'
  encodings_type = 'chexpert'
  out_path = '/deep/u/markendo/CXR-RePaiR/results/CheXpert/Retrieval-Baseline/generated_reports.csv'
  if encodings_type == 'chexpert':
      mimic_h5py_path = '/deep/group/report-clip/resnet_mimic_encodings/chexpert.h5'
      # out_path = '/deep/u/markendo/R2Gen/out/retrieval_baseline_mimic_generated_reports.csv'
      state_dict_path = '/deep/u/markendo/aihc-winter19-robustness/chexpert-model/classification_model_checkpoints/resnet18/1z6xfh2n/epoch=1-chexpert_competition_AUROC=0.88_v0.ckpt'
  elif encodings_type == 'moco-cxr':
      mimic_h5py_path = '/deep/group/report-clip/resnet_mimic_encodings/moco-cxr.h5'
      # out_path = '/deep/group/report-clip/aihc-win21-clip/notebooks/eval/out/moco_normalized/moco_normalized_generated_reports.csv'
      state_dict_path = '/deep/group/aihc-bootcamp-spring2020/cxr_fewer_samples/experiments/ntruongv/r8w1n416s_20201202_resnet18-chexpert-0.0001-pretrained_20201202-093038_SLURM1931462/checkpoint_0000.pth.tar'
  
  time_before = time.time()
  
  # load resnet 18 model
  model = models.resnet18(pretrained=False)
  model.fc = Identity()
  state_dict = torch.load(state_dict_path)
  state_dict = state_dict['state_dict']
  if encodings_type == 'chexpert':
      for key in list(state_dict.keys()):
          new_key = key.replace('model.model.', '')
          if (new_key != "fc.bias" and new_key != "fc.weight"):
              state_dict[new_key] = state_dict[key]
          del state_dict[key]
  elif encodings_type == 'moco-cxr':
      for key in list(state_dict.keys()):
          if 'encoder_q' in key: # discard encoder_k, only use query not key
              new_key = key.replace('module.encoder_q.', '')
              state_dict[new_key] = state_dict[key]
          del state_dict[key]
      del state_dict['fc.0.weight']
      del state_dict['fc.0.bias']
      del state_dict['fc.2.weight']
      del state_dict['fc.2.bias']
  model.load_state_dict(state_dict)
  model = torch.nn.DataParallel(model)
  model.to(device)
  
  if save_train_encodings:
      mimic_transform = transforms.Compose([
          # means computed from sample in `cxr_stats` notebook
          transforms.Normalize((101.48761, 101.48761, 101.48761), (83.43944, 83.43944, 83.43944)),
      ])
      mimic_dset = MIMICDataset(img_path=cxr_filepath,
                              txt_path=txt_filepath, transform=mimic_transform)
      mimic_loader = data.DataLoader(mimic_dset, batch_size=batch_size)
      # get the images and impressions
  
      if not os.path.exists(Path(mimic_h5py_path).parent.absolute()):
          os.makedirs(Path(mimic_h5py_path).parent.absolute())
  
      dset_size = len(mimic_loader.dataset)
      with h5py.File(mimic_h5py_path, 'w') as f:
          with torch.no_grad():
              encodings_dset = f.create_dataset('encodings', shape=(dset_size, 512))
              reports_dset = f.create_dataset('reports', shape=(dset_size,), dtype=h5py.string_dtype())
              for i, pack in enumerate(tqdm(mimic_loader)):
                  imgs = pack['img']
                  resnet_encodings = model(imgs)
                  encodings_arr = resnet_encodings.data.cpu().numpy()
                  reports = pack['txt']
  
                  start_index = i * batch_size
                  if i == len(mimic_loader) - 1:
                      encodings_dset[start_index:] = encodings_arr
                      reports_dset[start_index:] = reports
                  else:
                      encodings_dset[start_index:start_index+batch_size] = encodings_arr
                      reports_dset[start_index:start_index+batch_size] = reports
      saved_file = h5py.File(mimic_h5py_path)
      print(saved_file.get('encodings')[0])
      print(saved_file.get('reports')[0])
  
  
  # retrieve train examples that are most similar to train set encodings
  if retrieve_most_similar_train:
      # MIMIC-CXR test set
      # CXR_FILEPATH = '/deep/group/data/med-data/mimic_test_cxr.h5'
      # CXR_FILEPATH = '/deep/group/data/med-data/mimic-cxr-jpg-split/bootstrap_test/cxr.h5'
      # mimic_transform = transforms.Compose([
      #     # means computed from sample in `cxr_stats` notebook
      #     transforms.Normalize((101.48761, 101.48761, 101.48761), (83.43944, 83.43944, 83.43944)),
      # ])
      # # mimic_transform = transforms.Compose([
      # #             transforms.Resize((224, 224)),
      # #             transforms.ToTensor(),
      # #             transforms.Normalize((0.485, 0.456, 0.406),
      # #                                 (0.229, 0.224, 0.225))])
      # mimic_test_dset = MIMICTestDataset(CXR_FILEPATH, transform=mimic_transform)
      # loader = data.DataLoader(mimic_test_dset, shuffle=False, batch_size=batch_size)
  
  
      # CheXpert test set
      chexpert_transform = transforms.Compose([
          # transforms.ToTensor(),
          # transforms.Normalize((.5020, .5020, .5020),(.085585, .085585, .085585)),
          transforms.Normalize((129.4185, 129.4185, 129.4185), (73.3378, 73.3378, 73.3378))])
      torch_chexpert_dset = CheXpertDataset(img_path=chexpert_filepath, transform=chexpert_transform)
      chexpert_loader = data.DataLoader(torch_chexpert_dset, batch_size=batch_size)
  
      # data_mean = next(iter(chexpert_loader))
      # print(data_mean[0].mean(), data_mean[0].std())
  
      train_encodings_dset = MIMICEncodingsDataset(mimic_h5py_path)
      train_encodings_loader = data.DataLoader(train_encodings_dset, batch_size=256)
  
      output_reports = []
  
      with torch.no_grad():
          for i, images in enumerate(chexpert_loader):
              images = images.to(device)
              test_encodings_batch = model(images)
              batch_size = len(test_encodings_batch)
              highest_similarities = np.array([-1.] * batch_size)
              best_reports = [''] * batch_size
              for train_encodings in tqdm(train_encodings_loader):
                  train_encodings_batch = train_encodings[0]
  
                  # Using torch operations (for gpu)
                  train_encodings_batch = train_encodings_batch.to(device)
                  test_encodings_batch_norm = test_encodings_batch / test_encodings_batch.norm(dim=1)[:, None]
                  train_encodings_batch_norm = train_encodings_batch / train_encodings_batch.norm(dim=1)[:, None]
                  sim = torch.mm(test_encodings_batch_norm, train_encodings_batch_norm.transpose(0,1))
                  maxes = torch.max(sim, dim=1)
  
                  # Using np operations (for cpu)
                  # sim = cosine_similarity(test_encodings_batch.cpu(), train_encodings_batch)
                  # maxes = np.amax(sim, axis=1)
  
                  for minibatch_index in range(batch_size):
                      if maxes.values[minibatch_index] > highest_similarities[minibatch_index]:
                          highest_similarities[minibatch_index] = maxes.values[minibatch_index]
                          best_reports[minibatch_index] = train_encodings[1][maxes.indices[minibatch_index]].decode("utf-8")
                          # best_reports[minibatch_index] = train_encodings[1][np.argmax(sim[minibatch_index])].decode("utf-8")
              output_reports.extend(best_reports)
      print(time.time() - time_before)
      _df = pd.DataFrame(output_reports)
      _df.columns = ["report"]
      _df.to_csv(out_path, index=False)
apis: torch.utils.data.DataLoader torch.max torch.load pandas.read_csv pandas.DataFrame torch.no_grad numpy.repeat numpy.expand_dims torch.is_tensor torch.from_numpy torch.cuda.is_available numpy.shape numpy.array torch.nn.DataParallel