{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import numpy as np\n",
    "import os\n",
    "import shutil\n",
    "from tqdm import tqdm\n",
    "import os\n",
    "import shutil\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def retrieve_joint_connections(dataset):\n",
    "    if dataset == \"Human36M\":\n",
    "        joint_connections = [[14, 15], [15, 16], [13, 12], [12, 11], [9, 8], [8, 7], [4, 5], [5, 6], [3, 2], [2, 1], [7, 0], [0, 4], [0, 1], [8, 11], [8, 14], [9, 10]]    \n",
    "        defined_kpts = 17\n",
    "\n",
    "    elif dataset == \"face\":\n",
    "        joint_connections = [[0, 1], [1, 2], [2, 3], [3, 4],[5, 6], [6, 7], [7, 8], [8, 9], [10, 11], [11, 12], [12, 13], [14, 15], [15, 16], [16, 17], [17, 18], [19, 20], [20, 21], [21, 22], [22, 23], [23, 24], [24, 19], [25, 26], [26, 27], [27, 28], [28, 29], [29, 30], [30, 25], [31, 32], [32, 33], [33, 34], [34, 35], [35, 36], [36, 37], [37, 38], [38, 39], [39, 40], [40, 41], [41, 42], [42, 31], [43, 44], [44, 45], [45, 46], [46, 47], [47, 48], [48, 49], [49, 50], [50, 43]]\n",
    "        defined_kpts = 53\n",
    "\n",
    "    # elif dataset == \"cheetah\":\n",
    "    #     joint_connections = [[0,2], [0,1], [0,3], [2,3], [1,3], [3,4], [4,5], [5,6], [6,7],\n",
    "    #             [3,8], [8,9], [9,10], [10,11],\n",
    "    #             [3,12], [12,13], [13,14], [14,15],\n",
    "    #             [5, 16], [16,17], [17,18], [18,19],\n",
    "    #             [5,20], [20,21], [21,22], [22,23]]\n",
    "\n",
    "    elif dataset == \"cheetah\":\n",
    "        joint_connections = [[3, 1], [1, 2], [3, 4], [4, 5], \n",
    "                        [5, 6], [6, 7], [8, 9], [9, 10], [3, 8], \n",
    "                        [3, 11], [11, 12], [12, 13], [5, 17], \n",
    "                        [17, 18], [18, 19], [5, 14], [14, 15], \n",
    "                        [15, 16]]\n",
    "        defined_kpts = 20\n",
    "\n",
    "    elif dataset == \"cheetahtr\":\n",
    "        joint_connections = [[3, 1], [1, 2], [3, 4], [4, 5], \n",
    "                        [5, 6], [6, 7], [8, 9], [9, 10], [3, 8], \n",
    "                        [3, 11], [11, 12], [12, 13], [5, 17], \n",
    "                        [17, 18], [18, 19], [5, 14], [14, 15], \n",
    "                        [15, 16]]\n",
    "        defined_kpts = 20\n",
    "\n",
    "\n",
    "    elif dataset == \"cheetahsub\":\n",
    "        joint_connections = [[2, 0], [0, 1], [2, 3], [3, 4],\n",
    "                            [4, 5], [6, 7], [7, 8], [2, 6],\n",
    "                            [2, 9], [9, 10], [10, 11], [3, 15],\n",
    "                            [15, 16], [16, 17], [3, 12], [12, 13], [13, 14]]\n",
    "        defined_kpts = 18\n",
    "\n",
    "    elif dataset == \"hands\":\n",
    "        joint_connections = [[0,1],[1,2],[2,3],[3,4],[0,5],[5,6],[6,7],[7,8],[0,9],[9,10],[10,11],[11,12],[0,13],[13,14],[14,15],[15,16],[0,17],[17,18],[18,19],[19,20]]\n",
    "        defined_kpts = 21\n",
    "\n",
    "    elif dataset == \"amass\":\n",
    "        joint_connections = [[9, 13], [13, 16], [16, 18], [18, 20], [6,9], [3, 6], [0, 3], [0, 1], [0, 2], [1, 4] , [4, 7], [7, 10], [2, 5], [5, 8], [8, 11], \n",
    "                        [9, 14], [14, 17], [17, 19], [19, 21], [6, 12], [12, 15], \n",
    "                        [20, 34], [34, 35], [35, 36], \n",
    "                        [20, 22], [22, 23], [23, 24], \n",
    "                        [20, 25], [25, 26], [26, 27], \n",
    "                        [20, 31], [31, 32], [32, 33],\n",
    "                        [20, 28], [28, 29], [29, 30], \n",
    "                        #  [22, 25], [25, 31], [31, 28]\n",
    "                        [21, 49], [49, 50], [50, 51],\n",
    "                        [21, 37], [37, 38], [38, 39],\n",
    "                        [21, 40], [40, 41], [41, 42],\n",
    "                        [21, 46], [46, 47], [47, 48],\n",
    "                        [21, 43], [43, 44], [44, 45]]\n",
    "        defined_kpts = 52\n",
    "\n",
    "    elif dataset == \"openmonkey\":\n",
    "        joint_connections = [[0, 1], [1, 2], [2, 3], [3, 4], [2, 5], [5, 6], [2, 7], [7, 8], [8, 9], [7, 10], [10, 11], [7, 12]]    \n",
    "        defined_kpts = 13\n",
    "\n",
    "    elif dataset == \"wholebodyh36m\":\n",
    "        joint_connections = [[0, 1], [1, 3], [0, 2], [2, 4], [59, 64], [65, 70], [71, 82], \n",
    "                            [71, 83], [77, 87], [77, 88], [88, 89], [89, 90], [71, 90],\n",
    "                            [5, 7], [7, 9], [9, 91], [91, 92], [93, 96], [96, 100], [100, 104], [104, 108], [91, 108],\n",
    "                            [6, 8], [8, 10], [10, 112], [112, 113], [114, 117], [117, 121], [121, 125], [125, 129], [112, 129],\n",
    "                            [5, 6], [6, 12], [11, 12], [5, 11], [12, 14], [14, 16], [16, 20], [16, 21], [16, 22], [11, 13], [13, 15], [15, 17], [15, 18], [15, 19]]    \n",
    "\n",
    "        joint_connections = [\n",
    "            # Connect points as defined by the tuples\n",
    "            [0, 1], [1, 3], [0, 2], [2, 4], \n",
    "            [5, 7], [7, 9], [9, 91], #[91, 92], [93, 96], [96, 100], [100, 104], [104, 108], [91, 108],\n",
    "            [6, 8], [8, 10], [10, 112], #[112, 113], [114, 117], [117, 121], [121, 125], [125, 129], [112, 129],\n",
    "            [5, 6], [6, 12], [11, 12], [5, 11], \n",
    "            [12, 14], [14, 16], [16, 20], [16, 21], [16, 22], \n",
    "            [11, 13], [13, 15], [15, 17], [15, 18], [15, 19],\n",
    "\n",
    "            # Face\n",
    "            [23, 24], [24, 25], [25, 26], [26, 27], [27, 28], [28, 29], [29, 30], [30, 31], [31, 32], [32, 33], [33, 34],\n",
    "            [34, 35], [35, 36], [36, 37], [37, 38], [38, 39], [40, 41], [41, 42], [42, 43], [43, 44], [59, 60], [60, 61],\n",
    "            [61, 62], [62, 63], [63, 64], [59, 64], [45, 46], [46, 47], [47, 48], [48, 49], [65, 66], [66, 67], [67, 68],\n",
    "            [68, 69], [69, 70], [65, 70], [50, 51], [51, 52], [52, 53], [54, 55], [55, 56], [56, 57], [57, 58], [71, 72],\n",
    "            [72, 73], [73, 74], [74, 75], [75, 76], [76, 77], [77, 78], [78, 79], [79, 80], [80, 81], [81, 82], [82, 83],\n",
    "            [83, 84], [84, 85], [85, 86], [86, 87], [87, 88], [88, 89], [89, 90], [91, 92],\n",
    "\n",
    "            # Left hand\n",
    "            [91, 92], [92, 93], [93, 94], [94, 95], [91, 96], [96, 97], [97, 98], [98, 99], [91, 100], [100, 101], [101, 102],\n",
    "            [102, 103], [91, 104], [104, 105], [105, 106], [106, 107], [91, 108], [108, 109], [109, 110], [110, 111],\n",
    "\n",
    "\n",
    "            # Right hand\n",
    "            [112, 113], [113, 114], [114, 115], [115, 116], [112, 117], [117, 118], [118, 119], [119, 120], [112, 121],\n",
    "            [121, 122], [122, 123], [123, 124], [112, 125], [125, 126], [126, 127], [127, 128], [112, 129], [129, 130],\n",
    "            [130, 131], [131, 132]\n",
    "\n",
    "        ]    \n",
    "\n",
    "        defined_kpts = 133\n",
    "\n",
    "    elif dataset == \"bp4d+\":\n",
    "        joint_connections = [\n",
    "        # Left eyebrow (viewed from the model's perspective)\n",
    "        [15, 16], [16, 17], [17, 18], [18, 19],\n",
    "        [10, 11], [11, 12], [12, 13], [13, 14],\n",
    "        \n",
    "        # Right eyebrow\n",
    "        [0, 1], [1, 2], [2, 3], [3, 4],\n",
    "        [5, 6], [6, 7], [7, 8], [8, 9],\n",
    "        \n",
    "        # Bridge of the nose (from between the eyebrows to the tip)\n",
    "        [36, 37], [37, 38], [38, 39], [39, 40], [40, 41], [41, 42], [42, 43], [43, 44], [44, 45], [45, 46], [46, 47],\n",
    "        \n",
    "        # Left eye\n",
    "        [28, 29], [29, 30], [30, 31], [31, 32], [32, 33], [33, 34], [34, 35], [35, 28],\n",
    "        \n",
    "        # Right eye\n",
    "        [20, 21], [21, 22], [22, 23], [23, 24], [24, 25], [25, 26], [26, 27], [27, 20],\n",
    "        \n",
    "        # Outer part of the lips (outline of the lips)\n",
    "        [48, 49], [49, 50], [50, 51], [51, 52], [52, 53], [53, 54], [54, 55], [55, 56], [56, 57], [57, 58], [58, 59], [59, 48],\n",
    "        \n",
    "        # Inner part of the lips (detail within the lips)\n",
    "        [60, 61], [61, 62], [62, 63], [63, 64], [64, 65], [65, 66], [66, 67], [67, 60],\n",
    "        \n",
    "        # Jawline (from left ear, around the chin, to right ear)\n",
    "        [68, 69], [69, 70], [70, 71], [71, 72], [72, 73], [73, 74], [74, 75], [75, 76], [76, 77], [77, 78], [78, 79], [79, 80], [80, 81], [81, 82]\n",
    "    ]\n",
    "        defined_kpts = 83\n",
    "\n",
    "            \n",
    "    elif dataset == \"panoptic\":\n",
    "        joint_connections = [[0, 1], [0, 3], [3, 4], [4, 5], [0, 2], [2, 6], [6, 7], [7, 8], [2, 12], [12, 13], [13, 14], [0, 9], [9, 10], [10, 11]]\n",
    "        defined_kpts = 15\n",
    "\n",
    "    elif dataset == 'aeroplane':\n",
    "        joint_connections = [[2, 5], [1, 4], [5, 3], [3, 7], [7, 0], [0, 5], [5, 7], [5, 6], [6, 0], [6, 3], [2, 4], [2, 1]]\n",
    "        defined_kpts = 8\n",
    "\n",
    "    elif dataset == 'bicycle':\n",
    "        joint_connections = [[0,3], [0,7], [0, 2], [0, 6], [0, 10], [9, 10], [4, 10], [8, 10], [1, 9], [5, 9]]\n",
    "        defined_kpts = 11\n",
    "\n",
    "    elif dataset == \"tiger\" or dataset == \"cow\" or dataset == \"horse\" or dataset == \"hippo\" or dataset == \"dog\":\n",
    "        joint_connections = [[0,24], [0, 20], [1, 21], [1, 24], [7, 25], [19, 25], [6, 17],\n",
    "                        [4, 15], [3, 14], [9, 15], [8, 14], [9, 13], [8, 12],\n",
    "                        [2, 23], [2, 22], [2, 24], [11, 17], [10, 16], [5, 16],\n",
    "                        [7, 10], [7, 11], [13,18], [12, 18], [7, 18], [24,18]]\n",
    "        defined_kpts = 26\n",
    "\n",
    "    elif dataset == \"tigersubset\" or dataset == \"cowsubset\" or dataset == \"horsesubset\" or dataset == \"hipposubset\" or dataset == \"dogsubset\":\n",
    "        joint_connections = [[3, 17], [15, 17], [2, 13], [5, 9], [4, 8], [7, 13], [6, 12], [1, 12], [3, 6], [3, 7], [9, 14], [8, 14], [3, 14], [5, 11], [4, 10], [0, 14], [0, 16]] \n",
    "        defined_kpts = 18\n",
    "\n",
    "    elif dataset == 'boat':\n",
    "        joint_connections = [[0, 2], [0, 3], [0, 1], [1, 2], [1, 3], [2, 4], [3, 5], [4, 5], [1, 5], [1, 4]]\n",
    "        defined_kpts = 6\n",
    "\n",
    "    elif dataset == 'bottle':\n",
    "        joint_connections = [[0, 1], [1, 2], [0, 2], [3, 4], [3, 5], [4, 5], [1, 4], [0, 3], [2, 5], [1, 6], [0, 6], [2, 6]]\n",
    "        defined_kpts = 8\n",
    "\n",
    "    elif dataset == 'busfull' or dataset == 'bus':\n",
    "        joint_connections = [[5, 7], [4, 5], [6, 7], [4, 6], [1, 5], [1, 3], [3, 7], [0, 1], [2, 3], [0, 2], [2, 10], [0, 8], [8, 9], [10, 11], [6, 11], [4, 9]]\n",
    "        defined_kpts = 12\n",
    "\n",
    "    elif dataset == 'car':\n",
    "        joint_connections = [[0, 8], [0, 4], [4, 10], [8, 10],\n",
    "                                [10, 9], [9, 11], [8, 11], [11, 6], \n",
    "                                [9, 2], [2, 6], [4, 1], [5, 1], \n",
    "                                [0, 5], [5, 7], [1, 3], [7, 3], [3, 2], [7, 6]]    \n",
    "        defined_kpts = 12\n",
    "\n",
    "    elif dataset == 'busmissing':\n",
    "        joint_connections = [[5, 7], [4, 5], [6, 7], [4, 6], [1, 5], [1, 3], [3, 7], [0, 1], [2, 3], [0, 2], [2, 6], [0, 4]]\n",
    "        defined_kpts = 8\n",
    "\n",
    "    elif dataset == 'diningtable':\n",
    "        joint_connections = [[0, 2], [4, 6], [1, 3], [5, 7], [1, 5], [3, 7], [0, 4], [2, 6], [0, 1], [2, 3], [4, 5], [6, 7]]\n",
    "        defined_kpts = 8\n",
    "\n",
    "    elif dataset == 'tvmonitor':\n",
    "        joint_connections = [[5, 7], [4, 5], [4, 6], [6, 7], [0, 1], [0, 2], [2, 3], [1, 3], [3, 7], [1, 5], [2, 6], [0, 4]]\n",
    "        defined_kpts = 8\n",
    "\n",
    "    elif dataset == 'train':\n",
    "        joint_connections = [[4, 5], [4, 6], [6, 7], [5, 7], [0, 1], [1, 3], [2, 3], [0, 2], [1, 5], [0, 4], [2, 6], [3, 7], [1, 5]]\n",
    "        defined_kpts = 8\n",
    "\n",
    "\n",
    "    elif dataset == 'train16':\n",
    "        joint_connections = [[0, 1], [1,5], [5, 9], [9, 15], [3, 7], [7, 11], [11, 13], [2, 3], [2, 6], [6, 10], [10, 12], [1, 3], [0, 2], \n",
    "                            [0, 4], [4, 8], [8, 14], [15, 13], [13, 12], [12, 14], [14, 15]]\n",
    "        defined_kpts = 16\n",
    "\n",
    "    elif dataset == 'motorbike':\n",
    "        joint_connections = [[6, 2], [2, 9], [2, 3], [3, 8], [5, 8],\n",
    "                                [3, 5], [2, 1], [1, 0], [0, 7], [0, 4],\n",
    "                                [4, 7], [1, 4], [1, 7], [1, 5], [1, 8]]\n",
    "        defined_kpts = 10\n",
    "\n",
    "    elif dataset == 'sofa':\n",
    "        joint_connections = [[1, 5], [5, 4], [4, 6], [6, 2], [2, 0], \n",
    "                                [1, 0], [0, 4], [1, 3], [7, 5], [2, 3], \n",
    "                                [3, 7], [9, 7], [7, 6], [6, 8], [8, 9]]\n",
    "        defined_kpts = 10\n",
    "\n",
    "    elif dataset == 'chair':\n",
    "        joint_connections = [[7, 3], [6, 2], [9, 5], [8, 4], [7, 9], \n",
    "                                [8, 6], [6, 7], [9, 8], [9, 1], [8, 0], [1, 0]]\n",
    "        defined_kpts = 10\n",
    "\n",
    "    # MBW datasets\n",
    "    elif dataset == 'colobusmonkey':\n",
    "        joint_connections = [[0, 1], [1, 2], [2, 3], [3, 4], [1, 5], [5, 6], [6, 7], [1, 8], [8, 9], [9, 10], [10, 11], [11, 12], [9, 13], [13, 14], [14, 15]]\n",
    "        defined_kpts = 16\n",
    "\n",
    "    elif dataset == 'chimpanzee':\n",
    "        joint_connections = [[0, 1], [1, 2], [2, 3], [3, 4], [1, 5], [5, 6], [6, 7], [1, 8], [8, 9], [9, 10], [10, 11], [11, 12], [9, 13], [13, 14], [14, 15]]\n",
    "        defined_kpts = 16\n",
    "\n",
    "\n",
    "    elif dataset == 'tigerzoo':\n",
    "        joint_connections = [[0, 1], [1, 2], [2, 3], [3,4], [4,5], [1,6], [6,7], [1,8], [8,9], [3,10], [10,11], [3,12], [12,13]]\n",
    "        defined_kpts = 14\n",
    "\n",
    "    elif dataset == 'clownfish':\n",
    "        joint_connections = [[0, 1], [1, 2], [2, 3], [1, 4], [1, 5]]\n",
    "        defined_kpts = 6\n",
    "\n",
    "    elif dataset == 'fish':\n",
    "        joint_connections = [[0, 1], [1, 2], [2, 3], [1, 3], [3, 4], [4, 5], [5, 6], [6, 7], [5, 7], [5, 8], [8, 9], [9, 10], [8, 10], [10, 11], [11, 0]]\n",
    "        defined_kpts = 12\n",
    "\n",
    "    elif dataset == 'seahorse':\n",
    "        joint_connections = [[0, 1], [1, 2], [2,3], [1,3], [3,4], [4,5]]\n",
    "        defined_kpts = 6\n",
    "    \n",
    "    return joint_connections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load pickle file\n",
    "\n",
    "category_name = 'aeroplane'\n",
    "suffix = 'val'\n",
    "\n",
    "pickle_name = category_name + '_' + suffix\n",
    "\n",
    "## Load input data\n",
    "input_data_path = pickle_name + '.pkl'\n",
    "with open(input_data_path, 'rb') as f:\n",
    "    input_data = pickle.load(f)\n",
    "\n",
    "input_2d = input_data['W_GT']\n",
    "image_path = input_data['image_path']\n",
    "\n",
    "## Load predictions pickle file\n",
    "pred_data_path = category_name + '_3dlfm.pkl'\n",
    "with open(pred_data_path, 'rb') as f:\n",
    "    pred_data = pickle.load(f)\n",
    "labels_3d = pred_data['labels_3d']\n",
    "outputs_3d = pred_data['outputs_3d']\n",
    "\n",
    "\n",
    "joint_connections = retrieve_joint_connections(category_name)\n",
    "\n",
    "## Print the statistics\n",
    "print(\"Number of images: \", len(image_path))\n",
    "print(\"Input 2D shape: \", input_2d.shape)\n",
    "print(\"Labels 3D shape: \", labels_3d.shape)\n",
    "print(\"Outputs 3D shape: \", outputs_3d.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_3d_skeleton(predictions_3d, labels_3d, joint_connections, range_scale=2500, masks=None):\n",
    "    \"\"\"Visualize 3D skeletons for predicted and ground truth data.\"\"\"    \n",
    "\n",
    "    # Extract 3D coordinates and masks for the given sample index\n",
    "    pred_coordinates = predictions_3d\n",
    "    label_coordinates = labels_3d\n",
    "    \n",
    "    # Extract X, Y, Z coordinates after filtering\n",
    "    label_x, label_y, label_z = label_coordinates.T    \n",
    "\n",
    "    # Filter joint connections based on the mask\n",
    "    if masks is not None:\n",
    "        updated_connections = [connection for connection in joint_connections if masks[connection[0]] == 1.0 and masks[connection[1]] == 1.0]\n",
    "        print(\"Updated connections: {}\".format(updated_connections))\n",
    "    else:\n",
    "        updated_connections = joint_connections    \n",
    "\n",
    "    \n",
    "    \n",
    "    # Plotly Traces\n",
    "    traces = []\n",
    "    # Predicted skeleton\n",
    "    traces.extend(get_trace3d(updated_connections, pred_coordinates, 'blue', 'blue', \"Predicted KP\", masks=masks))\n",
    "    # Ground truth skeleton\n",
    "    traces.extend(get_trace3d(updated_connections, label_coordinates, 'red', 'red', \"Groundtruth KP\", masks=masks))\n",
    "\n",
    "    # Define layout\n",
    "    layout = go.Layout(\n",
    "        scene=dict(\n",
    "            aspectratio=dict(x=1, y=1, z=2),\n",
    "            xaxis=dict(range=[-label_x.max() * range_scale, label_x.max() * range_scale, ], showticklabels=False),\n",
    "            yaxis=dict(range=[-label_z.max() * range_scale, label_z.max() * range_scale], showticklabels=False),\n",
    "            zaxis=dict(range=[-label_y.max() * range_scale, label_y.max() * range_scale], showticklabels=False),\n",
    "        ),\n",
    "        width=700,\n",
    "        margin=dict(r=20, l=10, b=10, t=10),\n",
    "        scene_camera=dict(\n",
    "            up=dict(x=0, y=0, z=1),\n",
    "            center=dict(x=0, y=0, z=0),\n",
    "            eye=dict(x=0, y=-1.5, z=1.25),\n",
    "        )\n",
    "    )\n",
    "\n",
    "    # Create and display the plot\n",
    "    # fig = go.Figure(data=traces, layout=layout)\n",
    "    fig = go.Figure(data=traces)\n",
    "    fig.update_layout(scene=dict(aspectmode=\"data\"))        \n",
    "    fig.update_layout(\n",
    "        scene=dict(\n",
    "            xaxis=dict(title='', showticklabels=False),\n",
    "            yaxis=dict(title='', showticklabels=False),\n",
    "            zaxis=dict(title='', showticklabels=False)\n",
    "        )\n",
    "    )\n",
    "    fig.show()    \n",
    "\n",
    "\n",
    "def get_trace3d(joint_connections, points3d, point_color, line_color, name, masks=None):\n",
    "    \"\"\"Generate plotly traces for 3D points and connections.\"\"\"\n",
    "\n",
    "    # Filter 3D coordinates based on the mask\n",
    "    if masks is not None:\n",
    "        masked_coordinates = points3d[masks == 1.0]\n",
    "    else:\n",
    "        masked_coordinates = points3d\n",
    "    \n",
    "    x, z, y = masked_coordinates.T  # Swap Y and Z here\n",
    "    x_trace, z_trace, y_trace = points3d.T  # Swap Y and Z here\n",
    "\n",
    "    # Trace of points\n",
    "    trace_pts = go.Scatter3d(\n",
    "        x=x, y=y, z=z,\n",
    "        mode='markers',\n",
    "        name=name,\n",
    "        marker=dict(symbol='circle', size=6, color=point_color)\n",
    "    )\n",
    "\n",
    "    # Trace of lines\n",
    "    x_lines = []\n",
    "    y_lines = []\n",
    "    z_lines = []\n",
    "\n",
    "    for start, end in joint_connections:\n",
    "        x_lines.extend([x_trace[start], x_trace[end], None])\n",
    "        y_lines.extend([y_trace[start], y_trace[end], None])\n",
    "        z_lines.extend([z_trace[start], z_trace[end], None])\n",
    "\n",
    "    trace_lines = go.Scatter3d(\n",
    "        x=x_lines, y=y_lines, z=z_lines,\n",
    "        mode='lines',\n",
    "        name=name,\n",
    "        line=dict(width=6, color=line_color)\n",
    "    )\n",
    "\n",
    "    return [trace_pts, trace_lines]\n",
    "\n",
    "\n",
    "import plotly.graph_objects as go\n",
    "\n",
    "def plot_2d_skeleton(predictions_2d, labels_2d, joint_connections, masks=None):\n",
    "    \"\"\"Visualize 2D skeletons for predicted and ground truth data.\"\"\"\n",
    "    \n",
    "    # Extract 2D coordinates and masks for the given sample index\n",
    "    pred_coordinates = predictions_2d\n",
    "    label_coordinates = labels_2d\n",
    "\n",
    "    # Filter joint connections based on the mask\n",
    "    if masks is not None:\n",
    "        updated_connections = [connection for connection in joint_connections if masks[connection[0]] == 1.0 and masks[connection[1]] == 1.0]\n",
    "        print(\"Updated connections: {}\".format(updated_connections))\n",
    "    else:\n",
    "        updated_connections = joint_connections\n",
    "\n",
    "    \n",
    "    \n",
    "    \n",
    "    # Plotly Traces\n",
    "    traces = []\n",
    "    # Predicted skeleton\n",
    "    traces.extend(get_trace2d(updated_connections, pred_coordinates, 'blue', 'blue', \"Predicted KP\", masks=masks))\n",
    "    # Ground truth skeleton\n",
    "    traces.extend(get_trace2d(updated_connections, label_coordinates, 'red', 'red', \"Groundtruth KP\", masks=masks))\n",
    "\n",
    "    # Define layout\n",
    "    layout = go.Layout(\n",
    "        width=700,\n",
    "        height=700,\n",
    "        margin=dict(r=20, l=10, b=10, t=10)\n",
    "    )\n",
    "\n",
    "    # Create and display the plot\n",
    "    fig = go.Figure(data=traces, layout=layout)\n",
    "    fig.show()\n",
    "\n",
    "def get_trace2d(joint_connections, points2d, point_color, line_color, name, masks=None, get_lines=None):\n",
    "    \"\"\"Generate plotly traces for 2D points and connections.\"\"\"\n",
    "    \n",
    "    # Filter 2D coordinates based on the mask\n",
    "    if masks is not None:\n",
    "        masked_coordinates = points2d[masks == 1.0]\n",
    "    else:\n",
    "        masked_coordinates = points2d\n",
    "    \n",
    "    x, y = masked_coordinates.T  # Swap Y and Z here\n",
    "    x_trace, y_trace = points2d.T  # Swap Y and Z here\n",
    "\n",
    "    # Trace of points\n",
    "    trace_pts = go.Scatter(\n",
    "        x=x, y=y,\n",
    "        mode='markers',\n",
    "        name=name,\n",
    "        marker=dict(symbol='circle', size=6, color=point_color)\n",
    "    )\n",
    "\n",
    "    # Trace of lines\n",
    "    x_lines = []\n",
    "    y_lines = []\n",
    "\n",
    "    for start, end in joint_connections:\n",
    "        x_lines.extend([x_trace[start], x_trace[end], None])\n",
    "        y_lines.extend([y_trace[start], y_trace[end], None])\n",
    "\n",
    "    trace_lines = go.Scatter(\n",
    "        x=x_lines, y=y_lines,\n",
    "        mode='lines',\n",
    "        name=name,\n",
    "        line=dict(width=2, color=line_color)\n",
    "    )\n",
    "\n",
    "    if get_lines is not None:\n",
    "        if get_lines:\n",
    "            return [trace_pts, trace_lines]\n",
    "        else:\n",
    "            return [trace_pts]        \n",
    "    else:\n",
    "        return [trace_pts, trace_lines]\n",
    "\n",
    "\n",
    "import plotly.graph_objs as go\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "def plot_2d_skeleton_on_image(predictions_2d, labels_2d, joint_connections, image_path, masks=None, get_lines=None):\n",
    "    \"\"\"Visualize 2D skeletons for predicted and ground truth data on top of an image.\"\"\"\n",
    "    \n",
    "    # Load the image\n",
    "    image = Image.open(image_path)\n",
    "    width, height = image.size\n",
    "\n",
    "    # Extract 2D coordinates and masks for the given sample index\n",
    "    pred_coordinates = predictions_2d\n",
    "    label_coordinates = labels_2d\n",
    "    masks_ = masks\n",
    "    \n",
    "    # Filter joint connections based on the mask\n",
    "    if masks is not None:\n",
    "        updated_connections = [connection for connection in joint_connections if masks_[connection[0]] == 1.0 and masks_[connection[1]] == 1.0]\n",
    "    else:\n",
    "        updated_connections = joint_connections\n",
    "\n",
    "    print(\"updated connections: {}\".format(updated_connections))\n",
    "\n",
    "    # Plotly Traces\n",
    "    traces = []\n",
    "    # Image as background\n",
    "    traces.append(go.Scatter(\n",
    "        x=[0, width],\n",
    "        y=[0, height],\n",
    "        mode=\"markers\",\n",
    "        marker_opacity=0,\n",
    "        hoverinfo=\"none\",\n",
    "        showlegend=False\n",
    "    ))\n",
    "\n",
    "    # Predicted skeleton\n",
    "    traces.extend(get_trace2d(updated_connections, pred_coordinates, 'blue', 'blue', None, masks_, get_lines=get_lines))\n",
    "    # Ground truth skeleton\n",
    "    traces.extend(get_trace2d(updated_connections, label_coordinates, 'red', 'red', None, masks_, get_lines=get_lines))\n",
    "\n",
    "    # Define layout\n",
    "    layout = go.Layout(\n",
    "        width=width,\n",
    "        height=height,\n",
    "        xaxis=dict(showgrid=False, zeroline=False, showticklabels=False, range=[0, width]),\n",
    "        yaxis=dict(showgrid=False, zeroline=False, showticklabels=False, range=[height, 0], scaleanchor=\"x\"),\n",
    "        images=[go.layout.Image(source=image, xref=\"x\", yref=\"y\", x=0, y=0, sizex=width, sizey=height, sizing=\"stretch\", opacity=1.0, layer=\"below\")],\n",
    "        margin=dict(r=10, l=10, b=10, t=10),\n",
    "        hovermode=\"closest\",\n",
    "        showlegend=False,  # Hide legend\n",
    "    )\n",
    "    \n",
    "    # Create and display the plot\n",
    "    fig = go.Figure(data=traces, layout=layout)\n",
    "    fig.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "get_lines = False; frame_idx = 227\n",
    "plot_3d_skeleton(-outputs_3d[frame_idx], -labels_3d[frame_idx], joint_connections, range_scale=2500, masks=None)\n",
    "plot_2d_skeleton_on_image(input_2d[frame_idx], input_2d[frame_idx], joint_connections, image_path[frame_idx], masks=None, get_lines=get_lines)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Final storage #####\n",
    "\n",
    "## Randomly choose 10 frames and store them as a final pickle file\n",
    "# random_indices = random.sample(range(len(image_path)), 10)\n",
    "random_indices = [167, 29, 1, 3, 4, 7, 10, 18, 123, 227]\n",
    "final_data = {}\n",
    "final_data['image_path'] = np.asarray(image_path)[random_indices].tolist()\n",
    "final_data['inputs_2d'] = input_2d[random_indices]\n",
    "final_data['labels_3d'] = labels_3d[random_indices]\n",
    "final_data['outputs_3d'] = outputs_3d[random_indices]\n",
    "\n",
    "# Create final directory\n",
    "if not os.path.exists('final'):\n",
    "    os.makedirs('final')\n",
    "\n",
    "# if the below directory already exists, then delete it\n",
    "if os.path.exists('final/' + category_name + '_images'):\n",
    "    shutil.rmtree('final/' + category_name + '_images')\n",
    "os.makedirs('final/' + category_name + '_images', exist_ok=True)\n",
    "# Copying images to the new directory and updating the paths\n",
    "new_image_paths = []\n",
    "for path in tqdm(final_data['image_path']):\n",
    "    original_image_path = path # Save the original path\n",
    "    new_path = os.path.join('final/' + category_name + '_images', path)\n",
    "\n",
    "    # Create the directory if it doesn't exist\n",
    "    os.makedirs(os.path.dirname(new_path), exist_ok=True)\n",
    "\n",
    "    shutil.copy(original_image_path, new_path)\n",
    "    new_image_paths.append(new_path)\n",
    "\n",
    "# Update the image_path in data\n",
    "final_data['image_path'] = new_image_paths\n",
    "\n",
    "# Save the new pickle file\n",
    "final_pickle_name = 'final/' + category_name + '.pkl'\n",
    "with open(final_pickle_name, 'wb') as f:\n",
    "    pickle.dump(final_data, f)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(random_indices)\n",
    "random_indices = [167, 29, 1, 3, 4, 7, 10, 18, 123, 227]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "lifting",
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}