{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import numpy as np\n",
    "import os\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "data_path = # enter matterport data path\n",
    "uuid =  # enter UUID of the matterport model\n",
    "files = os.listdir(data_path)\n",
    "files = [file for file in files if file.startswith(uuid) and file.endswith('jpg')] # select all texture files\n",
    "grid_size = int(np.ceil(np.sqrt(len(files))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "target_texture_grid_size = 1024 # select a target texture size, the raw size is 2048, choose between 512,1024 \n",
    "# or 2048 depending on how many images you have\n",
    "\n",
    "\n",
    "combined_img = np.zeros((grid_size * target_texture_grid_size, grid_size * target_texture_grid_size, 3))\n",
    "files = sorted(files)\n",
    "#print(files)\n",
    "for num_file,file in enumerate(files):\n",
    "    img = np.array(Image.open(os.path.join(data_path, file)).resize((target_texture_grid_size,\n",
    "                                                                     target_texture_grid_size)))\n",
    "    i = grid_size - 1 - num_file // grid_size\n",
    "    j = num_file % grid_size\n",
    "    combined_img[i*target_texture_grid_size: (i+1)*target_texture_grid_size, \n",
    "                 j*target_texture_grid_size:(j+1)*target_texture_grid_size, :] = img\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(combined_img.astype(np.uint8))\n",
    "# visualize combined image, it is basically tiled version of original textures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Image.fromarray(combined_img.astype(np.uint8)).save(os.path.join(data_path, 'combined.jpg'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# change the obj file to use the combined texture, and rename to mesh_z_up\n",
    "\n",
    "vt_mtl_dict = {}\n",
    "\n",
    "with open('{}/{}.obj'.format(data_path, uuid)) as f:\n",
    "    for line in f:\n",
    "        if line.startswith('usemtl'):\n",
    "            current_mtl = int(line[-8:-5])\n",
    "            \n",
    "        if line.startswith('f'):\n",
    "            face = line.strip().split()\n",
    "            for vvt in face[1:]:\n",
    "                vt_mtl_dict[int(vvt.split('/')[1])] = current_mtl\n",
    "\n",
    "n_vt = 0\n",
    "fo = open('{}/{}.obj'.format(data_path, 'mesh_z_up'), 'w')\n",
    "mtl_set = False\n",
    "with open('{}/{}.obj'.format(data_path, uuid)) as f:\n",
    "    for line in f:\n",
    "        if line.startswith('usemtl'):\n",
    "            current_mtl = int(line[-8:-5])\n",
    "            #print(int(current_mtl))\n",
    "            if not mtl_set:\n",
    "                fo.write('usemtl combined\\n')\n",
    "            mtl_set = True\n",
    "        elif line.startswith('vt'):\n",
    "            n_vt += 1\n",
    "            n_mtl = vt_mtl_dict[n_vt]\n",
    "            i = n_mtl // grid_size\n",
    "            j = n_mtl % grid_size\n",
    "            _, u, v = line.strip().split()\n",
    "            u = float(u)\n",
    "            v = float(v)\n",
    "            new_u = u / grid_size + j / grid_size\n",
    "            new_v = v / grid_size + i / grid_size\n",
    "            #if n_mtl == 0:\n",
    "            fo.write('vt {} {}\\n'.format(new_u, new_v))\n",
    "            #else:\n",
    "            #    fo.write('vt {} {}\\n'.format(0, 0))\n",
    "        elif line.startswith('mtllib'):\n",
    "            fo.write('mtllib default.mtl\\n')\n",
    "        elif line.startswith('g'):\n",
    "            pass\n",
    "        else:\n",
    "            fo.write(line)\n",
    "            \n",
    "fo.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a mtl file\n",
    "with open(\"{}/default.mtl\".format(data_path), 'w') as fo:\n",
    "    fo.write(\n",
    "    \"\"\"\n",
    "newmtl combined\n",
    "Ns 0.000000\n",
    "Ka 1.000000 1.000000 1.000000\n",
    "Kd 0.000000 0.000000 0.000000\n",
    "Ks 1.000000 1.000000 1.000000\n",
    "Ke 0.000000 0.000000 0.000000\n",
    "Ni 1.450000\n",
    "d 1.000000\n",
    "illum 2\n",
    "map_Kd combined.jpg\n",
    "    \"\"\")\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# go to the folder and run `meshlab mesh_z_up.obj`, you should see the model being textured\n"
   ]
  },
  {
   "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": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (torch)",
   "language": "python",
   "name": "torch"
  },
  "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
