{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nerfstudio.cameras import camera_utils\n",
    "\n",
    "import torch\n",
    "import numpy as np\n",
    "import json\n",
    "import copy\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "project_name = 'site1_large'\n",
    "non_scaled_json = 'transforms.json'\n",
    "\n",
    "with open(os.path.join(project_name, non_scaled_json), 'r') as file1:\n",
    "    data1 = json.load(file1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.00841035273149741\n",
      "tensor([[ 9.9995e-01,  9.2652e-05, -9.6088e-03,  1.5047e-01],\n",
      "        [ 9.2652e-05,  9.9981e-01,  1.9283e-02,  6.8622e-02],\n",
      "        [ 9.6088e-03, -1.9283e-02,  9.9977e-01, -1.0862e-01]])\n"
     ]
    }
   ],
   "source": [
    "sorted_frames = sorted(data1[\"frames\"], key=lambda x: x[\"colmap_im_id\"])\n",
    "\n",
    "sorted_frames_raw = copy.deepcopy(sorted_frames)\n",
    "data1_raw = copy.deepcopy(data1)\n",
    "\n",
    "poses = []\n",
    "\n",
    "for frame in sorted_frames:\n",
    "    poses.append(np.array(frame[\"transform_matrix\"]))\n",
    "\n",
    "poses_scaled = torch.from_numpy(np.array(poses).astype(np.float32))\n",
    "\n",
    "poses_scaled, transform_matrix = camera_utils.auto_orient_and_center_poses(\n",
    "    poses_scaled,\n",
    "    method=\"up\",\n",
    "    center_method=\"poses\",\n",
    ")\n",
    "\n",
    "# print(poses_scaled)\n",
    "# print(transform_matrix)\n",
    "\n",
    "# Scale poses\n",
    "scale_factor = 1.0\n",
    "# auto_scale_poses:\n",
    "scale_factor /= float(torch.max(torch.abs(poses_scaled[:, :3, 3])))\n",
    "scale_factor *= 1.0\n",
    "\n",
    "poses_scaled[:, :3, 3] *= scale_factor\n",
    "\n",
    "print(scale_factor)\n",
    "print(transform_matrix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1893, 3, 4])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "poses_scaled.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "for idx,frame in enumerate(sorted_frames):\n",
    "    sorted_frames[idx][\"transform_matrix\"] = poses_scaled[idx].tolist()\n",
    "\n",
    "\n",
    "\n",
    "data1[\"frames\"] = sorted_frames\n",
    "data1_raw[\"frames\"] = sorted_frames_raw"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Writing JSON data\n",
    "# with open('normalized_transforms.json', 'w') as f:\n",
    "#     json.dump(data1, f, indent=4)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Validate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(4, 4)\n",
      "(4, 4)\n",
      "[[-0.14346903  0.09823143 -0.98476757 -0.56352723]\n",
      " [-0.98962168 -0.00609525  0.14356821  0.7554721 ]\n",
      " [ 0.00810051  0.99514493  0.09808644 -0.01924648]\n",
      " [ 0.          0.          0.          0.00841035]]\n",
      "[[-0.14346902072429657, 0.09823144227266312, -0.9847675561904907, -0.5635272264480591], [-0.989621639251709, -0.0060952454805374146, 0.143568217754364, 0.7554720044136047], [0.008100509643554688, 0.9951449632644653, 0.09808643162250519, -0.019246479496359825]]\n"
     ]
    }
   ],
   "source": [
    "i = 20\n",
    "\n",
    "matrix_1 = np.array(data1_raw['frames'][i][\"transform_matrix\"])\n",
    "matrix_2 = np.array(transform_matrix)\n",
    "\n",
    "new_row = np.array([[0, 0, 0, 1]])\n",
    "\n",
    "# matrix_1 = np.vstack((matrix_1, new_row))\n",
    "matrix_2 = np.vstack((matrix_2, new_row))\n",
    "\n",
    "print(matrix_1.shape)\n",
    "print(matrix_2.shape)\n",
    "\n",
    "result = np.dot(matrix_2, matrix_1)\n",
    "result[:, 3] *= scale_factor\n",
    "\n",
    "print(result)\n",
    "\n",
    "print(data1['frames'][i][\"transform_matrix\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Save Transform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "camera_norm_transform = transform_matrix.tolist()\n",
    "camera_norm_transform.append([0.0, 0.0, 0.0, 1.0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[0.9999538064002991,\n",
       "  9.265206608688459e-05,\n",
       "  -0.009608808904886246,\n",
       "  0.15047475695610046],\n",
       " [9.265206608688459e-05,\n",
       "  0.9998140931129456,\n",
       "  0.01928257942199707,\n",
       "  0.06862244009971619],\n",
       " [0.009608808904886246,\n",
       "  -0.01928257942199707,\n",
       "  0.9997678995132446,\n",
       "  -0.10861872136592865],\n",
       " [0.0, 0.0, 0.0, 1.0]]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "camera_norm_transform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the filename\n",
    "filename = os.path.join(project_name, \"transform_matrix.json\")\n",
    "\n",
    "# Step 1: Read the JSON file\n",
    "with open(filename, 'r') as f:\n",
    "    data = json.load(f)\n",
    "\n",
    "data[\"c2n_transform\"] = camera_norm_transform\n",
    "data[\"norm_scale\"] = scale_factor\n",
    "\n",
    "# Step 2: Write the flattened data back to the JSON file\n",
    "with open(filename, 'w') as f:\n",
    "    json.dump(data, f, indent=4)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nerfstudio2",
   "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.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}