prepare_ap10k.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "\n",
    "def load_data(*file_paths):\n",
    "    \"\"\"Load and merge data from multiple JSON files.\"\"\"\n",
    "    merged_data = {'annotations': [], 'images': [], 'categories': []}\n",
    "    for path in file_paths:\n",
    "        with open(path, 'r') as file:\n",
    "            data = json.load(file)\n",
    "            merged_data['annotations'].extend(data['annotations'])\n",
    "            merged_data['images'].extend(data['images'])\n",
    "            # Assuming categories are the same across files or only needed from the first file\n",
    "            if 'categories' in data and not merged_data['categories']:\n",
    "                merged_data['categories'] = data['categories']\n",
    "    return merged_data\n",
    "\n",
    "def remove_duplicate_annotations(data):\n",
    "    \"\"\"Remove duplicate annotations from the data based on image_id.\"\"\"\n",
    "    unique_image_ids = set()\n",
    "    new_annotations = []\n",
    "    for annotation in data['annotations']:\n",
    "        if annotation['image_id'] not in unique_image_ids:\n",
    "            unique_image_ids.add(annotation['image_id'])\n",
    "            new_annotations.append(annotation)\n",
    "    data['annotations'] = new_annotations\n",
    "    return data\n",
    "\n",
    "# Paths to your JSON files\n",
    "json_path_1 = \"data/ap-10k/annotations/ap10k-train-split1.json\"\n",
    "json_path_2 = \"data/ap-10k/annotations/ap10k-test-split1.json\"\n",
    "json_path_3 = \"data/ap-10k/annotations/ap10k-val-split1.json\"\n",
    "\n",
    "# Load and merge data\n",
    "data = load_data(json_path_1, json_path_2, json_path_3)\n",
    "\n",
    "# Remove duplicate annotations\n",
    "data = remove_duplicate_annotations(data)\n",
    "\n",
    "print(len(data['annotations']), \"unique annotations after cleanup.\")\n",
    "\n",
    "annotations = data[\"annotations\"]\n",
    "images = data[\"images\"]\n",
    "categories = data[\"categories\"]\n",
    "\n",
    "# Create dictionaries to index images by image_id and categories by category_id\n",
    "images_dict = {image[\"id\"]: image for image in images}\n",
    "categories_dict = {cat[\"id\"]: cat for cat in categories}\n",
    "\n",
    "# Set the base directory\n",
    "base_dir = \"data/ap-10k/ImageAnnotation\"\n",
    "\n",
    "# Function to pad the file name\n",
    "def pad_filename(filename):\n",
    "    return filename.zfill(17)\n",
    "\n",
    "new_annotations = []\n",
    "\n",
    "# Iterate over annotations to merge data and directly create and rename files\n",
    "for annotation in annotations:\n",
    "    image_id = annotation[\"image_id\"]\n",
    "    category_id = annotation.get(\"category_id\")\n",
    "\n",
    "    # Ensure both image_id and category_id exist in their respective dictionaries\n",
    "    if image_id in images_dict and category_id in categories_dict:\n",
    "        # Merge data from annotations, images, and categories\n",
    "        new_annotation = {**annotation, **images_dict[image_id],\n",
    "                          \"name\": categories_dict[category_id][\"name\"],\n",
    "                          \"supercategory\": categories_dict[category_id][\"supercategory\"]}\n",
    "        new_annotations.append(new_annotation)\n",
    "\n",
    "        supercategory = new_annotation.get(\"supercategory\")\n",
    "        name = new_annotation.get(\"name\")\n",
    "\n",
    "        # Check if supercategory and name are available\n",
    "        if supercategory and name:\n",
    "            # Create the directory path\n",
    "            directory_path = os.path.join(base_dir, supercategory, name)\n",
    "\n",
    "            # Create the directory if it doesn't exist\n",
    "            if not os.path.exists(directory_path):\n",
    "                os.makedirs(directory_path)\n",
    "\n",
    "            # Set the file name using the annotation id (or \"unknown\" if not available), and pad it\n",
    "            file_name = pad_filename(str(new_annotation.get(\"id\", \"unknown\")) + \".json\")\n",
    "            file_path = os.path.join(directory_path, file_name)\n",
    "\n",
    "            # Write the merged data to a JSON file\n",
    "            with open(file_path, 'w') as file:\n",
    "                json.dump(new_annotation, file, indent=4)\n",
    "\n",
    "print(\"All JSON files have been saved and renamed.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# split the images into different folders according to the supercategory and name\n",
    "\n",
    "import os\n",
    "import shutil\n",
    "\n",
    "# create a mapping dict to store the {supercategory}/{name} for each image id\n",
    "id_to_category = {}\n",
    "for item in new_annotations:\n",
    "    image_id = int(item.get(\"id\", \"unknown\"))\n",
    "    supercategory = item.get(\"supercategory\")\n",
    "    name = item.get(\"name\")\n",
    "    if supercategory and name:\n",
    "        id_to_category[image_id] = os.path.join(supercategory, name)\n",
    "\n",
    "# image folder path\n",
    "img_folder = \"data/ap-10k/data\"\n",
    "\n",
    "# iterate through the image folder\n",
    "for img_file in os.listdir(img_folder):\n",
    "    img_path = os.path.join(img_folder, img_file)\n",
    "    if os.path.isfile(img_path):\n",
    "        # extract the image id\n",
    "        img_id = int(img_file.split(\".\")[0][7:])\n",
    "        \n",
    "        # find the corresponding {supercategory}/{name}\n",
    "        category_path = id_to_category.get(img_id)\n",
    "        if category_path:\n",
    "            # create the target directory\n",
    "            target_directory = os.path.join(\"data/ap-10k/JPEGImages\", category_path)\n",
    "            if not os.path.exists(target_directory):\n",
    "                os.makedirs(target_directory)\n",
    "\n",
    "            # move the image to the target directory\n",
    "            target_path = os.path.join(target_directory, img_file)\n",
    "            shutil.move(img_path, target_path)\n",
    "\n",
    "# delete the original folder, data/ap-10k/data\n",
    "shutil.rmtree(\"data/ap-10k/data\")\n",
    "\n",
    "print(\"All images have been organized according to their respective categories.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "JSON file splitting complete.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import json\n",
    "\n",
    "base_path = \"data/ap-10k/ImageAnnotation\"\n",
    "\n",
    "def load_list_from_file(file_path):\n",
    "    \"\"\"Load a list of lines from a file.\"\"\"\n",
    "    with open(file_path, 'r') as file:\n",
    "        return [line.strip() for line in file]\n",
    "\n",
    "def split_json_files(json_files):\n",
    "    \"\"\"Split JSON files into train, eval, and test sets based on specified criteria.\"\"\"\n",
    "    length = len(json_files)\n",
    "    test_size = min(30, length)  # Ensure we don't exceed the total number of files\n",
    "    eval_size = min(20, max(0, length - test_size))  # Adjust eval_size based on remaining files\n",
    "    train_size = max(0, length - test_size - eval_size)  # Avoid negative train_size\n",
    "\n",
    "    # Split the dataset\n",
    "    train, eval, test = json_files[:train_size], json_files[train_size:train_size+eval_size], json_files[-test_size:]\n",
    "    return train, eval, test\n",
    "\n",
    "def save_list_to_file(file_path, items):\n",
    "    \"\"\"Save a list of items to a file.\"\"\"\n",
    "    with open(file_path, 'w') as file:\n",
    "        for item in items:\n",
    "            file.write(f\"{item}\\n\")\n",
    "\n",
    "is_crowd_list = load_list_from_file(\"data/ap-10k_is_crowd.txt\")\n",
    "\n",
    "for root, _, files in os.walk(base_path):\n",
    "    if root.count(os.sep) == base_path.count(os.sep) + 2:\n",
    "        json_list = [os.path.join(root, f) for f in files if f.endswith(\".json\")]\n",
    "        filtered_json_list = []\n",
    "\n",
    "        for json_file in json_list:\n",
    "            with open(json_file, 'r') as f:\n",
    "                json_data = json.load(f)\n",
    "            \n",
    "            if json_file.split('/')[-1].strip('.json') in is_crowd_list:\n",
    "                json_data[\"is_crowd\"] = 1\n",
    "                with open(json_file, 'w') as f:\n",
    "                    json.dump(json_data, f)\n",
    "            elif json_data[\"num_keypoints\"] >= 3:\n",
    "                filtered_json_list.append(json_file)\n",
    "        \n",
    "        train_json_list, eval_json_list, test_json_list = split_json_files(filtered_json_list)\n",
    "        \n",
    "        # Save the lists to their respective files\n",
    "        save_list_to_file(os.path.join(root, \"train_filtered.txt\"), train_json_list)\n",
    "        save_list_to_file(os.path.join(root, \"val_filtered.txt\"), eval_json_list)\n",
    "        save_list_to_file(os.path.join(root, \"test_filtered.txt\"), test_json_list)\n",
    "\n",
    "print(\"JSON file splitting complete.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "import json\n",
    "import numpy as np\n",
    "import itertools\n",
    "\n",
    "# Set the seed for reproducibility\n",
    "random.seed(42)\n",
    "\n",
    "def generate_pairs(base_path, file_name, output_folder, N_multiplier=None): #intra-specicies\n",
    "    N_total_pairs = 0\n",
    "    for root, dirs, files in os.walk(base_path):\n",
    "        # Process only subdirectories that are two levels deep\n",
    "        if root.count(os.sep) == base_path.count(os.sep) + 2:\n",
    "            json_list = []\n",
    "            with open(os.path.join(root, file_name), 'r') as file:\n",
    "                json_list = [line.strip() for line in file]\n",
    "\n",
    "            # For training, set N based on the length of json_list and a multiplier\n",
    "            if N_multiplier is not None:\n",
    "                N = 50 * len(json_list)  # Specific to training\n",
    "            else:\n",
    "                N = len(list(itertools.combinations(json_list, 2)))\n",
    "\n",
    "            random.shuffle(json_list)\n",
    "            all_possible_pairs = list(itertools.combinations(json_list, 2))\n",
    "            possible_pairs = []\n",
    "\n",
    "            for pair in all_possible_pairs:\n",
    "                src_json_path, trg_json_path = pair\n",
    "                with open(src_json_path) as f:\n",
    "                    src_data = json.load(f)\n",
    "                with open(trg_json_path) as f:\n",
    "                    trg_data = json.load(f)\n",
    "                src_kpt = np.array(src_data[\"keypoints\"]).reshape(-1, 3)\n",
    "                trg_kpt = np.array(trg_data[\"keypoints\"]).reshape(-1, 3)\n",
    "                mutual_vis = src_kpt[:, -1] / 2 * trg_kpt[:, -1] / 2\n",
    "                if mutual_vis.sum() >= 3:\n",
    "                    possible_pairs.append(pair)\n",
    "\n",
    "            # Adjust N based on the number of possible pairs\n",
    "            N = min(N, len(possible_pairs))\n",
    "            pairs_sampled = random.sample(possible_pairs, N) if N > 0 else []\n",
    "\n",
    "            for pair in pairs_sampled:\n",
    "                src_json_path, trg_json_path = pair\n",
    "                src_json_name, trg_json_name = os.path.basename(src_json_path).split(\".\")[0], os.path.basename(trg_json_path).split(\".\")[0]\n",
    "                category_name = os.path.basename(os.path.dirname(src_json_path))\n",
    "                new_json_file_name = f\"{src_json_name}-{trg_json_name}:{category_name}.json\"\n",
    "                \n",
    "                if not os.path.exists(output_folder):\n",
    "                    os.makedirs(output_folder)\n",
    "                new_json_path = os.path.join(output_folder, new_json_file_name)\n",
    "                \n",
    "                data = {\"src_json_path\": src_json_path, \"trg_json_path\": trg_json_path}\n",
    "                with open(new_json_path, 'w') as f:\n",
    "                    json.dump(data, f, indent=4)\n",
    "                \n",
    "                N_total_pairs += 1\n",
    "\n",
    "    print(f\"Total {N_total_pairs} pairs generated for {file_name}\")\n",
    "\n",
    "base_path = \"data/ap-10k/ImageAnnotation\"\n",
    "# Generate for training set with a specific multiplier for N\n",
    "generate_pairs(base_path, \"train_filtered.txt\", 'data/ap-10k/PairAnnotation/trn', N_multiplier=50)\n",
    "# Generate for test and validation sets without the multiplier for N\n",
    "generate_pairs(base_path, \"test_filtered.txt\", 'data/ap-10k/PairAnnotation/test')\n",
    "generate_pairs(base_path, \"val_filtered.txt\", 'data/ap-10k/PairAnnotation/val')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "import json\n",
    "import numpy as np\n",
    "import itertools\n",
    "\n",
    "# Set the seed for reproducibility\n",
    "random.seed(42)\n",
    "\n",
    "def generate_cross_species_pairs(base_path, file_name, output_folder, N_pairs_per_combination):\n",
    "    N_total_pairs = 0\n",
    "    subfolder_path = {}\n",
    "    for root, dirs, files in os.walk(base_path):\n",
    "        if root.count(os.sep) == base_path.count(os.sep) + 1:\n",
    "            subfolder_path[root] = []\n",
    "            for subroot, subdirs, subfiles in os.walk(root):\n",
    "                if subroot.count(os.sep) == root.count(os.sep) + 1:\n",
    "                    subfolder_path[root].append(subroot)\n",
    "\n",
    "    for key, value in subfolder_path.items():\n",
    "        if len(value) > 1:  # the family has more than one species\n",
    "            total_cross_species_pairs = []\n",
    "            species_combination = list(itertools.combinations(value, 2))\n",
    "            for species_pair in species_combination:\n",
    "                with open(os.path.join(species_pair[0], file_name), 'r') as train_file:\n",
    "                    train_json_list_1 = [line.strip() for line in train_file]\n",
    "\n",
    "                with open(os.path.join(species_pair[1], file_name), 'r') as train_file:\n",
    "                    train_json_list_2 = [line.strip() for line in train_file]\n",
    "\n",
    "                cross_species_pairs = list(itertools.product(train_json_list_1, train_json_list_2))\n",
    "                for pair in cross_species_pairs:\n",
    "                    if random.random() > 0.5:\n",
    "                        pair = (pair[1], pair[0])\n",
    "                total_cross_species_pairs.extend(cross_species_pairs)\n",
    "\n",
    "            possible_pairs = []\n",
    "            for pair in total_cross_species_pairs:\n",
    "                src_json_path, trg_json_path = pair\n",
    "                with open(src_json_path) as f:\n",
    "                    src_data = json.load(f)\n",
    "                with open(trg_json_path) as f:\n",
    "                    trg_data = json.load(f)\n",
    "                src_kpt = np.array(src_data[\"keypoints\"]).reshape(-1, 3)\n",
    "                trg_kpt = np.array(trg_data[\"keypoints\"]).reshape(-1, 3)\n",
    "                src_vis = src_kpt[:, -1] / 2\n",
    "                trg_vis = trg_kpt[:, -1] / 2\n",
    "                mutual_vis = src_vis * trg_vis\n",
    "                if mutual_vis.sum() >= 3:\n",
    "                    possible_pairs.append(pair)\n",
    "\n",
    "            N = min(N_pairs_per_combination, len(possible_pairs))\n",
    "            pairs_sampled = random.sample(possible_pairs, N)\n",
    "\n",
    "            for pair in pairs_sampled:\n",
    "                src_json_path, trg_json_path = pair\n",
    "                src_json_name = os.path.basename(src_json_path).split(\".\")[0]\n",
    "                trg_json_name = os.path.basename(trg_json_path).split(\".\")[0]\n",
    "                category_name = os.path.basename(os.path.dirname(os.path.dirname(src_json_path)))\n",
    "                new_json_file_name = f\"{src_json_name}-{trg_json_name}:{category_name}.json\"\n",
    "                new_json_path = os.path.join(output_folder, new_json_file_name)\n",
    "\n",
    "                if not os.path.exists(output_folder):\n",
    "                    os.makedirs(output_folder)\n",
    "\n",
    "                data = {\n",
    "                    \"src_json_path\": src_json_path,\n",
    "                    \"trg_json_path\": trg_json_path\n",
    "                }\n",
    "\n",
    "                with open(new_json_path, 'w') as f:\n",
    "                    json.dump(data, f, indent=4)\n",
    "\n",
    "            N_total_pairs += N\n",
    "    \n",
    "    print(f\"Total {N_total_pairs} pairs for {file_name}\")\n",
    "\n",
    "base_path = \"data/ap-10k/ImageAnnotation\"\n",
    "generate_cross_species_pairs(base_path, \"val_filtered.txt\", 'data/ap-10k/PairAnnotation/val_cross_species', 400)\n",
    "generate_cross_species_pairs(base_path, \"test_filtered.txt\", 'data/ap-10k/PairAnnotation/test_cross_species', 900)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "import json\n",
    "import numpy as np\n",
    "import itertools\n",
    "from tqdm import tqdm\n",
    "\n",
    "# Set seed for reproducibility\n",
    "random.seed(42)\n",
    "\n",
    "def generate_cross_family_pairs(base_path, file_name, output_folder, N_pairs_per_combination):\n",
    "    N_total_pairs = 0\n",
    "    subfolder_path = {}\n",
    "\n",
    "    # Collect all JSON file paths within each family\n",
    "    for root, dirs, files in os.walk(base_path):\n",
    "        if root.count(os.sep) == base_path.count(os.sep) + 1:\n",
    "            subfolder_path[root] = []\n",
    "            for subroot, subdirs, subfiles in os.walk(root):\n",
    "                if subroot.count(os.sep) == root.count(os.sep) + 1:\n",
    "                    with open(os.path.join(subroot, file_name), 'r') as train_file:\n",
    "                        train_json_list = [line.strip() for line in train_file]\n",
    "                    subfolder_path[root].extend(train_json_list)\n",
    "\n",
    "    # Iterate over combinations of families\n",
    "    families_combination = list(itertools.combinations(subfolder_path.keys(), 2))\n",
    "\n",
    "    for family_pair in tqdm(families_combination, desc=\"Generating pairs\"):\n",
    "        family_1, family_2 = family_pair\n",
    "        family_1_json_list = subfolder_path[family_1]\n",
    "        family_2_json_list = subfolder_path[family_2]\n",
    "\n",
    "        # Generate all possible pairs between two families\n",
    "        cross_family_pairs = list(itertools.product(family_1_json_list, family_2_json_list))\n",
    "\n",
    "        # Filter out pairs based on visibility criteria\n",
    "        possible_pairs = []\n",
    "        for pair in cross_family_pairs:\n",
    "            src_json_path, trg_json_path = pair\n",
    "            with open(src_json_path) as f:\n",
    "                src_data = json.load(f)\n",
    "            with open(trg_json_path) as f:\n",
    "                trg_data = json.load(f)\n",
    "            src_kpt = np.array(src_data[\"keypoints\"]).reshape(-1, 3)\n",
    "            trg_kpt = np.array(trg_data[\"keypoints\"]).reshape(-1, 3)\n",
    "            mutual_vis = src_kpt[:, -1] / 2 * trg_kpt[:, -1] / 2\n",
    "            if mutual_vis.sum() >= 3:\n",
    "                possible_pairs.append(pair)\n",
    "\n",
    "        # Sample pairs if more are available than needed\n",
    "        N = min(N_pairs_per_combination, len(possible_pairs))\n",
    "        pairs_sampled = random.sample(possible_pairs, N)\n",
    "\n",
    "        # Create annotation files for sampled pairs\n",
    "        for pair in pairs_sampled:\n",
    "            src_json_path, trg_json_path = pair\n",
    "            src_json_name, trg_json_name = os.path.basename(src_json_path).split(\".\")[0], os.path.basename(trg_json_path).split(\".\")[0]\n",
    "            category_name = os.path.basename(os.path.dirname(os.path.dirname(src_json_path)))\n",
    "            new_json_file_name = f\"{src_json_name}-{trg_json_name}:all.json\"\n",
    "            new_json_path = os.path.join(output_folder, new_json_file_name)\n",
    "\n",
    "            if not os.path.exists(output_folder):\n",
    "                os.makedirs(output_folder)\n",
    "\n",
    "            data = {\"src_json_path\": src_json_path, \"trg_json_path\": trg_json_path}\n",
    "            with open(new_json_path, 'w') as f:\n",
    "                json.dump(data, f, indent=4)\n",
    "\n",
    "            N_total_pairs += N\n",
    "\n",
    "    print(f\"Total {N_total_pairs} pairs generated for {file_name}\")\n",
    "\n",
    "base_path = \"data/ap-10k/ImageAnnotation\"\n",
    "# Generate for test set\n",
    "generate_cross_family_pairs(base_path, \"test_filtered.txt\", 'data/ap-10k/PairAnnotation/test_cross_family', 30)\n",
    "# Generate for val set\n",
    "generate_cross_family_pairs(base_path, \"val_filtered.txt\", 'data/ap-10k/PairAnnotation/val_cross_family', 20)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize the key points of ap-10k benchmark\n",
    "\n",
    "import json\n",
    "import numpy as np\n",
    "import torch\n",
    "from PIL import Image\n",
    "from glob import glob\n",
    "from utils.utils_correspondence import resize\n",
    "from utils.utils_visualization import draw_correspondences_gathered\n",
    "from utils.utils_dataset import preprocess_kps_pad\n",
    "\n",
    "np.random.seed(42)\n",
    "\n",
    "def load_ap10k_data(path=\"data/ap-10k\", size=840, category='all', split='test_cross_family', subsample=20):\n",
    "    np.random.seed(42)\n",
    "    pairs = sorted(glob(f'{path}/PairAnnotation/{split}/*:{category}.json'))\n",
    "    if subsample is not None and subsample > 0:\n",
    "        pairs = [pairs[ix] for ix in np.random.choice(len(pairs), subsample)]\n",
    "    # print(f'Number of SPairs for {category} = {len(pairs)}')\n",
    "    files = []\n",
    "    # print(f'Number of SPair key points for {category} <= {num_kps}')\n",
    "    kps = []\n",
    "    thresholds = []\n",
    "    for pair in pairs:\n",
    "        with open(pair) as f:\n",
    "            data = json.load(f)\n",
    "        source_json_path = data[\"src_json_path\"]\n",
    "        target_json_path = data[\"trg_json_path\"]\n",
    "        src_img_path = source_json_path.replace(\"json\", \"jpg\").replace('ImageAnnotation', 'JPEGImages')\n",
    "        trg_img_path = target_json_path.replace(\"json\", \"jpg\").replace('ImageAnnotation', 'JPEGImages')\n",
    "\n",
    "        with open(source_json_path) as f:\n",
    "            src_file = json.load(f)\n",
    "        with open(target_json_path) as f:\n",
    "            trg_file = json.load(f)\n",
    "            \n",
    "        source_bbox = np.asarray(src_file[\"bbox\"])  # l t w h\n",
    "        target_bbox = np.asarray(trg_file[\"bbox\"])\n",
    "        \n",
    "        source_size = np.array([src_file[\"width\"], src_file[\"height\"]])  # (W, H)\n",
    "        target_size = np.array([trg_file[\"width\"], trg_file[\"height\"]])  # (W, H)\n",
    "\n",
    "        # print(source_raw_kps.shape)\n",
    "        source_kps = torch.tensor(src_file[\"keypoints\"]).view(-1, 3).float()\n",
    "        source_kps[:,-1] /= 2\n",
    "        source_kps, src_x, src_y, src_scale = preprocess_kps_pad(source_kps, source_size[0], source_size[1], size)\n",
    "\n",
    "        target_kps = torch.tensor(trg_file[\"keypoints\"]).view(-1, 3).float()\n",
    "        target_kps[:,-1] /= 2\n",
    "        target_kps, trg_x, trg_y, trg_scale = preprocess_kps_pad(target_kps, target_size[0], target_size[1], size)\n",
    "        # The source thresholds aren't actually used to evaluate PCK on SPair-71K, but for completeness\n",
    "        # they are computed as well:\n",
    "        # thresholds.append(max(source_bbox[3] - source_bbox[1], source_bbox[2] - source_bbox[0]))\n",
    "        if split == 'test':\n",
    "            thresholds.append(max(target_bbox[3], target_bbox[2])*trg_scale)\n",
    "        elif split == 'trn':\n",
    "            thresholds.append(max(source_bbox[3], source_bbox[2])*src_scale)\n",
    "            thresholds.append(max(target_bbox[3], target_bbox[2])*trg_scale)\n",
    "\n",
    "        kps.append(source_kps)\n",
    "        kps.append(target_kps)\n",
    "        files.append(src_img_path)\n",
    "        files.append(trg_img_path)\n",
    "\n",
    "    kps = torch.stack(kps)\n",
    "    used_kps, = torch.where(kps[:, :, 2].any(dim=0))\n",
    "    kps = kps[:, used_kps, :]\n",
    "    # print(f'Final number of used key points: {kps.size(1)}')\n",
    "    return files, kps, thresholds, used_kps\n",
    "\n",
    "\n",
    "files, kps, thresholds, used_kps = load_ap10k_data(category='all', split='test_cross_family') # one can specify the category and split\n",
    "\n",
    "pair_idx = 0\n",
    "src_img = Image.open(files[2*pair_idx]).convert('RGB')\n",
    "src_kps = kps[2*pair_idx]\n",
    "src_img = resize(src_img, 840, resize=True, to_pil=True)\n",
    "\n",
    "trg_img = Image.open(files[2*pair_idx+1]).convert('RGB')\n",
    "trg_kps = kps[2*pair_idx+1]\n",
    "trg_img = resize(trg_img, 840, resize=True, to_pil=True)\n",
    "\n",
    "vis = src_kps[:, 2] * trg_kps[:, 2] > 0 #tensor([ True,  True,  True,  True,  True,  True,  True,  True, False,  True, True, False, False, False, False])\n",
    "src_kps = src_kps[vis]\n",
    "trg_kps = trg_kps[vis]\n",
    "fig = draw_correspondences_gathered(src_kps[:, [1,0]], trg_kps[:, [1,0]], src_img, trg_img)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sd-dino",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}