code: def main(args, config):
      # TODO - add the low shot part (one example of each?)
      # TODO - use configuration instead
      num_neighbors = 10
      metric_type = "cosine"
      normalize_features = True
      cosine_temperature = 0.1
      output_path = "/checkpoint/qduval/temp/"
  
      # Load the dataset
      train_batch_size = config.DATA.TRAIN.BATCHSIZE_PER_REPLICA
      train_dataset = load_dataset(config, split="train")
      train_data_loader = data.DataLoader(
          train_dataset, batch_size=train_batch_size, shuffle=False
      )
      val_batch_size = config.DATA.TEST.BATCHSIZE_PER_REPLICA
      val_dataset = load_dataset(config, split="val")
      val_data_loader = data.DataLoader(
          val_dataset, batch_size=val_batch_size, shuffle=False
      )
  
      # Build the model and load the weights from the checkpoint
      model = build_model(config["MODEL"], config["OPTIMIZER"]).cuda()
      checkpoint = CheckpointLoader.load_and_broadcast_init_weights(
          checkpoint_path=config.MODEL.WEIGHTS_INIT.PARAMS_FILE,
          device=torch.device("cpu"),
      )
      model.init_model_from_weights_params_file(config, checkpoint)
  
      # Iterate over the dataset and build the a table of pairs
      # (representation of patch, corresponding dominant class)
  
      print("Extracting train embeddings...")
      train_embeddings, train_targets, _ = extract_embeddings(
          model,
          train_data_loader,
          normalize_features=normalize_features,
      )
  
      print("Extracting val embeddings...")
      (
          val_embeddings,
          val_targets,
          val_full_targets,
      ) = extract_embeddings(
          model, val_data_loader, normalize_features=normalize_features
      )
  
      print("Building the KNN index...")
      train_keys = train_embeddings.numpy()
      if metric_type == "cosine":
          index = faiss.IndexFlatIP(train_embeddings.shape[-1])
      else:
          index = faiss.IndexFlatL2(train_embeddings.shape[-1])
      index.add(train_keys)
  
      print("Computing predictions for validation set...")
      val_keys = val_embeddings.numpy()
  
      if metric_type == "l2":
          assert num_neighbors == 1, "L2 only support 1 neighbor"
          distances, indices = index.search(val_keys, 1)
          predictions = train_targets[indices].squeeze(-1)
      else:
          assert num_neighbors > 0
          distances, indices = index.search(val_keys, num_neighbors)
          retrieved_neighbors = (
              train_targets[indices].view(-1, 1).long()
          )
          num_examples = val_embeddings.shape[0]
          num_classes = val_targets.max().item() + 1
          retrieval_one_hot = torch.zeros(
              num_examples * num_neighbors, num_classes
          )
          retrieval_one_hot.scatter_(
              1, retrieved_neighbors.view(-1, 1), 1
          )
          predictions = get_sorted_predictions(
              batch_size=num_examples,
              num_classes=num_classes,
              distances=torch.from_numpy(distances),
              retrieval_one_hot=retrieval_one_hot,
              temperature=cosine_temperature,
          )
          predictions = predictions.narrow(
              dim=1, start=0, length=1
          )  # Top-1
  
      print(
          predictions.shape, val_targets.shape, val_full_targets.shape
      )
  
      print("Saving the predictions...")
      # TODO - do not hard-code the size of the patches
      np.save(
          f"{output_path}/val_predictions.npy",
          predictions.reshape(-1, 14, 14),
      )
      np.save(
          f"{output_path}/val_targets.npy",
          val_targets.reshape(-1, 14, 14),
      )
  
      print("Computing metrics...")
      predictions = predictions.numpy()
      val_targets = val_targets.numpy()
      accuracy = sklearn.metrics.accuracy_score(
          predictions, val_targets
      )
      jaccard_score = sklearn.metrics.jaccard_score(
          predictions, val_targets, average=None
      )
      jaccard_score_weighted = sklearn.metrics.jaccard_score(
          predictions, val_targets, average="weighted"
      )
      average_jaccard_score = np.mean(jaccard_score)
      print(f"Accuracy: {accuracy}")
      print(f"Jaccard Score: {jaccard_score}")
      print(f"Jaccard Score (Mean): {average_jaccard_score}")
      print(f"Jaccard Score (Mean Weighted): {jaccard_score_weighted}")
  
      print("Computing full resolution metrics...")
      # TODO - do not hard-code the size of the patches
      high_res_predictions = F.interpolate(
          torch.from_numpy(predictions)
          .reshape((-1, 1, 14, 14))
          .float(),
          size=(val_full_targets.shape[-2], val_full_targets.shape[-1]),
          mode="bilinear",
      ).long()
      jaccard_score_weighted = sklearn.metrics.jaccard_score(
          high_res_predictions.reshape((-1,)).numpy(),
          val_full_targets.reshape((-1,)).numpy(),
          average="weighted",
      )
      print(f"Jaccard Score (Mean Weighted): {jaccard_score_weighted}")
  
      print("Saving the full resolution predictions...")
      np.save(
          f"{output_path}/val_full_predictions.npy",
          high_res_predictions.squeeze(1).numpy(),
      )
      np.save(f"{output_path}/val_full_targets.npy", val_full_targets)
  
      print("Computing full resolution binary metrics...")
      jaccard_score_weighted = sklearn.metrics.jaccard_score(
          high_res_predictions.reshape((-1,)).numpy() != 0,
          val_full_targets.reshape((-1,)).numpy() != 0,
          average="weighted",
      )
      print(f"Jaccard Score (Mean Weighted): {jaccard_score_weighted}")
  
file_path: vissl/tools/knn_segmentation.py
imports: import sys
  from typing import Any, List
  import faiss
  import numpy
  import sklearn.metrics
  import torch
  import torch.nn.functional
  import torch.utils.data
  import torchvision.datasets
  import torchvision.transforms
  from tqdm import tqdm
  from vissl.data import dataset_catalog
  from vissl.models import build_model
  from vissl.utils.checkpoint import CheckpointLoader
  from vissl.utils.hydra_config import compose_hydra_configuration, convert_to_attrdict
  from vissl.utils.knn_utils import get_sorted_predictions
label_desc: Use image node_id 3 for calling a global function with example usage: main(args, config) without return types