open_ClipMatrix_ver0_8.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wOwF0l7TRoA3"
      },
      "source": [
        "# Introduction to ClipMatrix: text controlled 3D mesh deformation and stylization.\n",
        "\n",
        "A novel type of text controlled AI generative art: you give a text and get a 3D textured shape from it. \n",
        "This code is loosely based on my recent [paper](https://arxiv.org/abs/2109.12922), but has many more tricks inside. Feel free to play with it, make art or even better: experiment with new ML techniques based on it.\n",
        "Please cite my twitter account [@NJetchev](https://twitter.com/NJetchev) or the paper when using this code, posting artwork from it or modifying it.\n",
        "\n",
        "Inputs from you:\n",
        "- starting mesh template\n",
        "- text prompt to be interpreted by CLIP\n",
        "\n",
        "Outputs: videos of the resulting generated mesh. by default saved in **clip_videos** folder\n",
        "\n",
        "Cool features:\n",
        "- check my account for artwork and inspiration what is possible [@NJetchev](https://twitter.com/NJetchev) \n",
        "- support for arbitrary template meshes for starting shapes\n",
        "- define regions of interest for fine part control\n",
        "- shadow mapping\n",
        "- normal mapping\n",
        "- [Sobolev smoothing](https://github.com/rgl-epfl/large-steps-pytorch) \n",
        "- neural shaders for advanced animation effects\n",
        "- you can export the 3D mesh file and print your work in a 3D printer service\n",
        "\n",
        "Libraries and references used: quite many, but I am most thankful to \n",
        "- OpenAI's CLIP \n",
        "- Meta's Pytorch3D\n",
        "- the authors of this [paper](https://rgl.epfl.ch/publications/Nicolet2021Large)\n",
        "- the researchers of [SMPL](https://smpl.is.tue.mpg.de/) , the sample static model.obj is taken from their website for research purposes\n",
        "- [@nonlethalcode](https://twitter.com/nonlethalcode) for helpful suggestions and testing of the UI of this code\n",
        "\n",
        "\n",
        "What is my plan for next updates of Open_ClipMatrix:\n",
        "- this code is related to the recent [ClipMatrix Creatures](https://clipmatrix.wordpress.com/) art project. If you appreciate it, holders of NFT Creatures will get earlier exclusive access to new versions of the ClipMatrix tool\n",
        "-  currently there are no riggable models inside, just static templates. I may add them in the future, not certain yet. It can complicate the code and increase runtimes a lot. \n",
        "- In this version I focused on a small and agile demo that is fun and fast to play with.\n",
        "\n",
        "**CHANGELOG:**\n",
        "- 0.65 initial version\n",
        "\n",
        "- 0.651 few bug fixed, options to disable shaders and normal maps\n",
        "\n",
        "- 0.7 added graph convolutions, fixed library paths; default regularization strength settings favor stronger mesh deformation\n",
        "\n",
        "- 0.8 better shadows, fixed library paths, support multiple template clones"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Setup\n"
      ],
      "metadata": {
        "id": "fk8itNyeScda"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "#os.environ['CUDA_LAUNCH_BLOCKING'] = \"1\"#for debug\n",
        "#os.environ['PYTORCH_CUDA_ALLOC_CONF'] = \"max_split_size_mb:32\"#for debug"
      ],
      "metadata": {
        "id": "Yg-Ds1uSKo5I"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Check GPU Status\n",
        "!nvidia-smi -L\n",
        "\n"
      ],
      "metadata": {
        "id": "UZ614ZnegkP_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title  Prepare Folders\n",
        "\n",
        "try:\n",
        "    from google.colab import drive\n",
        "    print(\"Google Colab detected. Using Google Drive.\")\n",
        "    is_colab = True\n",
        "    #@markdown If you connect your Google Drive, you can save the final outputs of each run on your drive.\n",
        "    google_drive = True #@param {type:\"boolean\"}\n",
        "except:\n",
        "    is_colab = False\n",
        "    google_drive = False\n",
        "    save_models_to_google_drive = False\n",
        "    print(\"Google Colab not detected.\")\n",
        "\n",
        "if is_colab:\n",
        "    if google_drive is True:\n",
        "        drive.mount('/content/drive')\n",
        "        resourcePath = '/content/drive/MyDrive/CLIPMatrix'\n",
        "    else:\n",
        "        resourcePath = '.'\n",
        "else:\n",
        "    resourcePath = '.'\n",
        "\n",
        "import os\n",
        "from os import path\n",
        "#Simple create paths taken with modifications from Datamosh's Batch VQGAN+CLIP notebook\n",
        "def createPath(filepath):\n",
        "    if path.exists(filepath) == False:\n",
        "      os.makedirs(filepath)\n",
        "      print(f'Made {filepath}')\n",
        "    else:\n",
        "      print(f'filepath {filepath} exists.')\n",
        "\n",
        "initDirPath = f'{resourcePath}/init_obj'\n",
        "createPath(initDirPath)\n",
        "outDirPath = f'{resourcePath}/output'\n",
        "createPath(outDirPath)\n",
        "videoDirPath = f'{outDirPath}/video'\n",
        "createPath(videoDirPath)\n",
        "meshDirPath = f'{outDirPath}/mesh'\n",
        "createPath(meshDirPath)\n",
        "\n",
        "#videos saved there\n",
        "vid_dir = videoDirPath\n",
        "#save video frames theres\n",
        "!mkdir frames2\n"
      ],
      "metadata": {
        "id": "tSHRiWttg-z1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown # [ 1.3 ] **download needed dependencies**\n",
        "#try:\n",
        "#  import cupy as cp\n",
        "#except:\n",
        "#  print (\"trying to import\")\n",
        "#  !pip install cupy-wheel\n",
        "#  import cupy as cp\n",
        "!pip install --no-deps ftfy regex tqdm\n",
        "!git clone https://github.com/openai/CLIP.git\n",
        "!apt install libsuitesparse-dev\n",
        "!pip install scikit-sparse\n",
        "#!git clone https://github.com/rgl-epfl/large-steps-pytorch.git\n",
        "#import sys\n",
        "#sys.path.append(\"large-steps-pytorch/largesteps\")\n",
        "!pip install largesteps"
      ],
      "metadata": {
        "id": "tjL2FJ45lIdX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xMnZGbP7WlcU"
      },
      "outputs": [],
      "source": [
        "#@markdown # [ 1.4 ] **Import Libraries 📚**\n",
        "import sys\n",
        "import os\n",
        "#can slow down os.environ['PYTORCH_CUDA_ALLOC_CONF'] = \"max_split_size_mb:100\"\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import matplotlib.pyplot as plt\n",
        "from skimage.io import imread\n",
        "import numpy as np\n",
        "import shutil\n",
        "import torchvision"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HGrnGjOQRoA_"
      },
      "source": [
        "## Instructions: settings to tune by User\n",
        "Note: this version has settings set to relatively low levels, so that it works out of the box even for weaker GPUs like p100. If you have a good machine with 16 GB or more GPU memory, set them to higher values, e.g. texture size, mesh subdivision, etc.\n",
        "\n",
        "You can decide how powerful and detailed the mesh you evolve will be, and how long you will compute it before saving the output video"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nB-pxfKsRoBC"
      },
      "outputs": [],
      "source": [
        "#@title Generator Options\n",
        "\n",
        "TEXSIZE = 512#@param {type:'raw'}#size in pixels of square texture image - larger sizes are costly\n",
        "HRL=1#@param {type:\"slider\", min:0, max:2, step:1} \n",
        "NClones=4#@param {type:\"slider\", min:1, max:10, step:1} \n",
        "\n",
        "train_num_views =3#@param {type:\"slider\", min:1, max:6, step:1} \n",
        "\n",
        "\n",
        "#if 0 use original mesh vertex resolution, if 1 or more - use mesh subdivision for more detail - more resource hungry\n",
        "mesh_obj_name = \"model.obj\"#@param {type:'raw'}#a sample human mesh for quick testing - feel free to try any 3d mesh you have\n",
        "# e.g. https://free3d.com/3d-models/ has many free resources\n",
        "\n",
        "sideX=768#@param {type:'raw'}#size of image for final videos we will be saving\n",
        "\n",
        "#prompts to use, we will sample randomly from them if multiple\n",
        "headprompt = \"head with 8 eyes \"#@param {type:'string'}, \"\"\n",
        "bodyprompt = \"torso armored body of \"#@param {type:'string'}, \"\"\n",
        "legprompt = \"legs and tail of \"#@param {type:'string'}, \"\"\n",
        "\n",
        "creatures = [\"Mecha Warrior\"] #@param {type:'raw'}\n",
        "detail = [\"detailed to the maximum\",\"detailed to the maximum\",\"rendered in Unreal Engine\"]#@param {type:'raw'}, \"\"\n",
        "styles = [\"in sci fi Giger style \"]#@param {type:'raw'}\n",
        "#note: you can also use an image as style -- simply add the image file as prompt.\n",
        "# e.g. download some \"image.jpg\" and set styles +=[\"image.jpg\"]\n",
        "\n",
        "#the larger you set these, the wilder the model will become and deform away from initial mesh\n",
        "initGamma=5e-2#@param {type:'raw'}\n",
        "initPV=0.01#@param {type:'raw'}\n",
        "\n",
        "optisteps = 296 #@param {type:\"slider\", min:100, max:600, step:1} \n",
        "#how many gradient steps for optimization\n",
        "\n",
        "extra_patch_clip =  1 #@param {type:\"slider\", min:1, max:9, step:1} \n",
        "clip_flip =  False  #@param {type:\"boolean\"}\n",
        "# how many extra patches to crop in CLIP, in addition to 3d cameras - costly\n",
        "\n",
        "fSobolevStrength = .2+4. #@param {type:'raw'}# the larger this is, the more regularized the mesh will be\n",
        "fPenaltyDeformation = 1e-1#@param {type:'raw'}#the larger this is , the smoother the mesh and close to initial one\n",
        "fPenaltyDeformation_e = 1e-2#@param {type:'raw'}\n",
        "fPartRegularize = 1e-2#@param {type:'raw'} # for clones if more than 1\n",
        "\n",
        "bWeightedEdges=False  #@param {type:\"boolean\"} \n",
        "#if True will use mesh edge distance for vertex Laplacian regularization\n",
        "viewport_limit=  49#@param {type:\"slider\", min:0, max:60, step:1} \n",
        "# how many degrees to sample for rotation views\n",
        "\n",
        "#regularize % of 3d mesh as part of total image - experimental\n",
        "fSilhTarget = 0.24#@param {type:'raw'}\n",
        "fSilhStrength = 4e1#@param {type:'raw'}\n",
        "\n",
        "#some neural renderer settings\n",
        "fShader_reg = 1e-2 #@param {type:'raw'}# with larger penalty, will be closed to underlying 3d model; with smaller penalty- let it be more wild\n",
        "channels_neural_shader = 16 #@param {type:\"slider\", min:0, max:80, step:1} \n",
        "# the more channels, the more memmory hungry but fancier shading you get\n",
        "\n",
        "useSymmetry = True #@param {type:\"boolean\"}\n",
        "#use x-axis mesh template symmetry\n",
        "useShadows = True #@param {type:\"boolean\"}\n",
        "#use shadows as effect in rendering\n",
        "\n",
        "##note: set these two options to false for better results if you are exporting mesh/texture to Blender\n",
        "useNormalMaps = True #@param {type:\"boolean\"}\n",
        "#use normal mapping for more detail\n",
        "useNeuralShader = True #@param {type:\"boolean\"}\n",
        "#use neural shading for more light effects\n",
        "\n",
        "bSOBOLEV = True #@param {type:\"boolean\"}\n",
        "\n",
        "create_video = True  #@param {type:\"boolean\"}\n",
        "bsave_obj = True #@param {type:\"boolean\"}\n",
        "save_texture = True #@param {type:\"boolean\"}\n",
        "show_3d_interactive_runall = True #@param {type:\"boolean\"}\n",
        "\n",
        "graphConv=True #@param {type:\"boolean\"}\n",
        "nGraphLayers = 4 #@param {type:\"slider\", min:2, max:9, step:1} \n",
        "\n",
        "fgraphnet=0.6\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ROAXQbTnSWtB"
      },
      "source": [
        "# Define Neccesary Functions\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OSfhRNYpf7ZX"
      },
      "outputs": [],
      "source": [
        "#torch3d\n",
        "need_pytorch3d=False\n",
        "try:\n",
        "    import pytorch3d\n",
        "except ModuleNotFoundError:\n",
        "    need_pytorch3d=True#works, needs restart however Nov. 11 2021    \n",
        "    #!python3 -c \"import torch;assert torch.__version__.startswith('1.6'), 'should be 1.6.x'\" || pip install torch==1.6.0+cu101 torchvision==0.7.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html --upgrade \n",
        "    #!pip install pytorch3d\n",
        "\n",
        "if need_pytorch3d:\n",
        "    print (\"needs pytorch\",torch.__version__,torch.version.cuda)\n",
        "    if False:#slow but works with latest COLAB, TODO switch to wheels once Pytorch3d fixed \n",
        "          !pip install 'git+https://github.com/facebookresearch/pytorch3d.git'\n",
        "    elif True:\n",
        "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
        "        version_str=\"\".join([f\"py3{sys.version_info.minor}_cu\",\\\n",
        "        torch.version.cuda.replace(\".\",\"\"),f\"_pyt{pyt_version_str}\"])\n",
        "        !pip install pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
        "    elif torch.__version__.startswith(\"1.9\") and sys.platform.startswith(\"linux\"):\n",
        "        # We try to install PyTorch3D via a released wheel.\n",
        "        version_str=\"\".join([\n",
        "            f\"py3{sys.version_info.minor}_cu\",\n",
        "            torch.version.cuda.replace(\".\",\"\"),\n",
        "            f\"_pyt{torch.__version__[0:5:2]}\"\n",
        "        ])\n",
        "        !pip install pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
        "    else:\n",
        "        # We try to install PyTorch3D from source.\n",
        "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
        "        !tar xzf 1.10.0.tar.gz\n",
        "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\n",
        "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'\n",
        "import pytorch3d\n",
        "print (pytorch3d.__version__)\n",
        "\n",
        "\n",
        "\n",
        "# Util function for loading meshes\n",
        "from pytorch3d.io import load_objs_as_meshes, save_obj\n",
        "\n",
        "from pytorch3d.loss import (\n",
        "    mesh_edge_loss, \n",
        "    mesh_laplacian_smoothing, \n",
        "    mesh_normal_consistency,\n",
        ")\n",
        "\n",
        "# Data structures and functions for rendering\n",
        "from pytorch3d.structures import Meshes\n",
        "from pytorch3d.renderer import (\n",
        "    look_at_view_transform,\n",
        "    OpenGLPerspectiveCameras, \n",
        "    PointLights, \n",
        "    DirectionalLights, \n",
        "    Materials, \n",
        "    RasterizationSettings, \n",
        "    MeshRenderer, \n",
        "    MeshRasterizer,  \n",
        "    SoftPhongShader,\n",
        "    SoftSilhouetteShader,\n",
        "    SoftPhongShader,\n",
        "    TexturesVertex\n",
        ")\n",
        "\n",
        "from pytorch3d.io import save_obj\n",
        "\n",
        "\n",
        "# add path for demo utils functions \n",
        "import sys\n",
        "import os\n",
        "sys.path.append(os.path.abspath(''))\n",
        "\n",
        "from pytorch3d.io import load_objs_as_meshes\n",
        "from PIL import Image\n",
        "from pytorch3d.io import load_obj\n",
        "from pytorch3d.ops import laplacian,norm_laplacian\n",
        "from pytorch3d.transforms import RotateAxisAngle\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ypI794oRxAKL"
      },
      "source": [
        "## Initial camera and large render sideX\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a4uOzSQEjckC"
      },
      "outputs": [],
      "source": [
        "# Setup\n",
        "if torch.cuda.is_available():\n",
        "    device = torch.device(\"cuda:0\")\n",
        "    torch.cuda.set_device(device)\n",
        "else:\n",
        "    device = torch.device(\"cpu\")\n",
        "\n",
        "# Place a point light in front of the object. As mentioned above, the front of the cow is facing the -z direction. \n",
        "canonicL= PointLights(device=device, location=[[0,0,10]])# PointLights(device=device, location=[[0.0, 0.0, -3.0]]) #\n",
        "lights = canonicL\n",
        "R, T = look_at_view_transform(dist=1.9, elev=20, azim=0)\n",
        "# We arbitrarily choose one particular view that will be used to visualize \n",
        "camera = OpenGLPerspectiveCameras(device=device, R=R,T=T) \n",
        "raster_settings = RasterizationSettings(\n",
        "    image_size=sideX, \n",
        "    blur_radius=0.0, \n",
        "    faces_per_pixel=1, perspective_correct=True\n",
        ")\n",
        "\n",
        "# Create a phong renderer by composing a rasterizer and a shader. The textured \n",
        "# phong shader will interpolate the texture uv coordinates for each vertex, \n",
        "# sample from a texture image and apply the Phong lighting model\n",
        "renderer = MeshRenderer(\n",
        "    rasterizer=MeshRasterizer(\n",
        "        cameras=camera, \n",
        "        raster_settings=raster_settings\n",
        "    ),\n",
        "    shader=SoftPhongShader(\n",
        "        device=device, \n",
        "        cameras=camera,\n",
        "        lights=lights\n",
        "    )\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LkSds3OOgXQT"
      },
      "outputs": [],
      "source": [
        "# Data structures and functions for rendering\n",
        "from pytorch3d.structures import Meshes\n",
        "from pytorch3d.vis.plotly_vis import AxisArgs, plot_batch_individually, plot_scene\n",
        "from pytorch3d.vis.texture_vis import texturesuv_image_matplotlib\n",
        "\n",
        "\n",
        "from pytorch3d.loss import (\n",
        "    mesh_edge_loss, \n",
        "    mesh_laplacian_smoothing, \n",
        "    mesh_normal_consistency,\n",
        ")\n",
        "\n",
        "from pytorch3d.renderer import (\n",
        "    look_at_view_transform,\n",
        "    FoVPerspectiveCameras, \n",
        "    FoVOrthographicCameras,\n",
        "    PointLights, \n",
        "    DirectionalLights, \n",
        "    Materials, \n",
        "    RasterizationSettings, \n",
        "    MeshRenderer, \n",
        "    MeshRasterizer,  \n",
        "    SoftPhongShader,\n",
        "    HardFlatShader,\n",
        "    TexturesUV,\n",
        "    TexturesVertex\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SR3etg15RoBQ"
      },
      "outputs": [],
      "source": [
        "from typing import Tuple, List\n",
        "from typing import Optional, Dict, Union\n",
        "from pytorch3d.renderer.utils import TensorProperties\n",
        "from typing import Optional,Union\n",
        "from pytorch3d.renderer.mesh.rasterizer import Fragments\n",
        "from pytorch3d.ops import interpolate_face_attributes\n",
        "from pytorch3d.renderer.mesh.shading import _apply_lighting\n",
        "from pytorch3d.renderer.blending import softmax_rgb_blend\n",
        "from pytorch3d.renderer.blending import BlendParams\n",
        "from pytorch3d.renderer.utils import TensorProperties\n",
        "from pytorch3d.renderer import SoftPhongShader"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "__sC90r3hkIS"
      },
      "outputs": [],
      "source": [
        "class deform_oracle(object):\n",
        "        def __init__(self):\n",
        "          self.net=None\n",
        "          self.last=None\n",
        "          self.shaderI=None\n",
        "        \n",
        "          self.collision=0\n",
        "\n",
        "        def calc(self,v_posed,pose_f=None,betas=None):\n",
        "          if self.last is not None:\n",
        "            return  v_posed+self.last\n",
        "\n",
        "          if self.net is None:\n",
        "            return 0\n",
        "          return self.net(v_posed,pose_f)\n",
        "\n",
        "ORACLE = deform_oracle()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qwYW-Hc2XIED"
      },
      "source": [
        "## UV textures"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!wget --no-check-certificate 'https://docs.google.com/uc?export=download&id=1CK2wrRqu94kPy3j4lm70UWaojrjgqO1F' -O model.obj"
      ],
      "metadata": {
        "id": "NNI6fi4Dfal7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from pytorch3d.ops.subdivide_meshes import SubdivideMeshes\n",
        "\n",
        "def subdivide(m,HRL=HRL):\n",
        "  if HRL==1:\n",
        "    return SubdivideMeshes()(m)\n",
        "  if HRL==2:\n",
        "    def SubSub(m):\n",
        "      s=SubdivideMeshes()\n",
        "      return s(s(m))\n",
        "    return SubSub(m)\n"
      ],
      "metadata": {
        "id": "rdT3Ulan6gMr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#all tensors 3dim-except faces_iuvs\n",
        "def repM(v10,f10,vt0,faces_uvs0):\n",
        "  if NClones ==1:\n",
        "    return v10,f10,vt0,faces_uvs0\n",
        "  bv=[]\n",
        "  bf=[]\n",
        "  bvt=[]\n",
        "  bft=[]\n",
        "\n",
        "  nv=v10.shape[1]\n",
        "  nvt=vt0.shape[1]\n",
        "  N=NClones\n",
        "  for i in range(N):\n",
        "    off=0\n",
        "    if i>0:\n",
        "      off=torch.zeros(1,1,3).cuda()#.uniform_(-0.0002,0.0002)\n",
        "    bv.append(v10+off)#hack, just some offset first\n",
        "    bf.append(f10+i*nv)\n",
        "    vt0_=vt0*1\n",
        "    #vt0_[...,0]/=N#shrink\n",
        "    #vt0_[...,0]+= i/N#offset\n",
        "    bvt.append(vt0_)#TODO other texture region\n",
        "    bft.append(faces_uvs0+i*nvt)\n",
        "    plt.scatter(vt0_[0,:,0].cpu(),vt0_[0,:,1].cpu())\n",
        "\n",
        "  plt.show()\n",
        "  return torch.cat(bv,1),torch.cat(bf,1),torch.cat(bvt,1),torch.cat(bft,0)\n",
        "\n"
      ],
      "metadata": {
        "id": "Cy6JwAkJDr5R"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cHn3Fgj0wGwt"
      },
      "outputs": [],
      "source": [
        "verts, faces, aux=load_obj(mesh_obj_name)\n",
        "\n",
        "v10=verts.unsqueeze(0).cuda()\n",
        "f10 = faces.verts_idx.unsqueeze(0).cuda()\n",
        "vt=aux.verts_uvs.cuda().unsqueeze(0)\n",
        "print(\"vt\",vt.max(),vt.min())\n",
        "faces_uvs = faces.textures_idx.cuda()\n",
        "print (vt.shape,faces_uvs.shape,faces_uvs.max())\n",
        "\n",
        "vt0 = vt*1\n",
        "faces_uvs0 = faces_uvs*1\n",
        "  \n",
        "  #raise Exception\n",
        "\n",
        "if HRL>0:##TODO add subdivide logic\n",
        "        \n",
        "    newm = subdivide(Meshes(v10,f10))\n",
        "    orig_v_template=newm.verts_padded().squeeze()\n",
        "    NVOrig = orig_v_template.shape[0]\n",
        "    print (\"NVorig\",NVOrig)\n",
        "    \n",
        "    f2=newm.faces_padded()\n",
        "    print(\"v2\",orig_v_template.shape)\n",
        "    print (\"faces\",f2.shape,f2.max(),\"old faces\",f10.shape,f10.max())\n",
        "    print (f2.max().item()+1)\n",
        "\n",
        "    #for multiscale - indices of clones low res mesh\n",
        "    ix_level0=[]\n",
        "    for i in range(NClones):#add low level for each clone\n",
        "      off=i*orig_v_template.shape[0]#verts in next cloned mesh, offset for initial level\n",
        "      ix_level0+=list(range(off,off+NVOrig))\n",
        "    _,f10,_,_=repM(v10,f10,vt0,faces_uvs0)#only faces used in multiscale logic       \n",
        "    \n",
        "    newm = subdivide(Meshes(torch.cat([vt0,vt0[:,:,:1]*0],2),faces_uvs0.unsqueeze(0)))\n",
        "    vt=newm.verts_padded()[...,:2]\n",
        "    faces_uvs=newm.faces_padded().squeeze()\n",
        "\n",
        "    ##repeat the cloning - -with high res mesh!\n",
        "    orig_v_template,f2,vt,faces_uvs=repM(orig_v_template.unsqueeze(0),f2,vt,faces_uvs)\n",
        "    orig_v_template=orig_v_template.squeeze()\n",
        "else:\n",
        "    #clone\n",
        "    print (v10.shape,f10.shape,vt0.shape,faces_uvs0.shape)\n",
        "    v10,f10,vt0,faces_uvs0=repM(v10,f10,vt0,faces_uvs0)\n",
        "    print (v10.shape,f10.shape,vt0.shape,faces_uvs0.shape)   \n",
        "    #other standard code \n",
        "    orig_v_template= v10.squeeze()\n",
        "    NVOrig = orig_v_template.shape[0]\n",
        "    f2=f10\n",
        "    faces_uvs=faces_uvs0\n",
        "    vt=vt0\n",
        "\n",
        "#raise Exception"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "if False:\n",
        "  fig = plot_scene({\n",
        "        \"text_to_add\": {\n",
        "            \"mean optimized\": m_check\n",
        "        }\n",
        "    })\n",
        "  fig.show()"
      ],
      "metadata": {
        "id": "PUeoQyfKWLjN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## define additional subregions of mesh - parts"
      ],
      "metadata": {
        "id": "OL5fYDmBLisC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# define anything appropriate for your template mesh\n",
        "# inspect manually\n",
        "\n",
        "iHead= orig_v_template[:,1]>0.16\n",
        "iTorso= (orig_v_template[:,1]>-0.46)*(orig_v_template[:,1]<0.16)*(orig_v_template[:,0].abs()<0.16)\n",
        "\n",
        "\n",
        "##parts only in mesh0\n",
        "print (iHead[:10],iHead[-10:])\n",
        "if NVOrig< iHead.shape[0]:# and False:\n",
        "  iHead[NVOrig:]=False\n",
        "  iTorso[NVOrig:]=False\n",
        "print (iHead[:10],iHead[-10:])\n",
        "\n",
        "plt.scatter(orig_v_template[:,0].cpu(),orig_v_template[:,1].cpu(),c=iHead.cpu())\n",
        "plt.show()\n",
        "plt.scatter(orig_v_template[:,0].cpu(),orig_v_template[:,2].cpu(),c=iHead.cpu())\n",
        "plt.show()\n",
        "\n",
        "\n",
        "plt.scatter(orig_v_template[:,0].cpu(),orig_v_template[:,1].cpu(),c=iTorso.cpu())\n",
        "plt.show()\n",
        "plt.scatter(orig_v_template[:,0].cpu(),orig_v_template[:,2].cpu(),c=iTorso.cpu())\n",
        "plt.show()\n",
        "\n",
        "\n",
        "parts = [iHead,iTorso]"
      ],
      "metadata": {
        "id": "xWb_Co39LmAt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "orig_v_template.shape"
      ],
      "metadata": {
        "id": "lnMj5XchOIFR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "iHead= orig_v_template[:,1]>0.16\n",
        "iTorso= (orig_v_template[:,1]>-0.4)*(orig_v_template[:,1]<0.16)*(orig_v_template[:,0].abs()<0.16)\n",
        "iLeg= (orig_v_template[:,1]<-0.46)\n",
        "\n",
        "if NVOrig< iHead.shape[0]:# and False:\n",
        "  iHead[NVOrig:]=False\n",
        "  iTorso[NVOrig:]=False\n",
        "  iLeg[NVOrig:]=False\n",
        "\n",
        "plt.scatter(orig_v_template[:,0].cpu(),orig_v_template[:,1].cpu(),c=iHead.cpu(),alpha=0.1)\n",
        "plt.show()\n",
        "plt.scatter(orig_v_template[:,0].cpu(),orig_v_template[:,2].cpu(),c=iHead.cpu(),alpha=0.1)\n",
        "plt.show()\n",
        "\n",
        "\n",
        "plt.scatter(orig_v_template[:,0].cpu(),orig_v_template[:,1].cpu(),c=iTorso.cpu(),alpha=0.1)\n",
        "plt.show()\n",
        "plt.scatter(orig_v_template[:,0].cpu(),orig_v_template[:,2].cpu(),c=iTorso.cpu(),alpha=0.1)\n",
        "plt.show()\n",
        "\n",
        "\n",
        "plt.scatter(orig_v_template[:,0].cpu(),orig_v_template[:,1].cpu(),c=iLeg.cpu(),alpha=0.1)\n",
        "plt.show()\n",
        "plt.scatter(orig_v_template[:,0].cpu(),orig_v_template[:,2].cpu(),c=iLeg.cpu(),alpha=0.1)\n",
        "plt.show()\n",
        "\n",
        "\n",
        "parts = [iHead,iTorso,iLeg]"
      ],
      "metadata": {
        "id": "eutGoENfapMV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z0FcgLvjyw8s"
      },
      "outputs": [],
      "source": [
        "#@title text prompts logic - aligned with parts definitions\n",
        "def sample(a):\n",
        "    return a[np.random.randint(len(a))].replace('\\n','')\n",
        "    \n",
        "def getTexts():\n",
        "    acreature= sample(creatures)\n",
        "    astyle= sample(styles)    \n",
        "    det =  detail[np.random.randint(len(detail))]\n",
        "    text_0 =   acreature + \" \" + astyle + \" \" + det\n",
        "    text_0bw = acreature+ \" \" + astyle\n",
        "\n",
        "    text_1  = str(headprompt)+\" \"+text_0\n",
        "    text_1bw = str(headprompt)+\" \"+text_0bw\n",
        "\n",
        "    text_2  = str(bodyprompt)+\" \"+text_0\n",
        "    text_2bw = str(bodyprompt)+\" \"+text_0bw \n",
        "\n",
        "    text_3 = str(legprompt) + text_0 \n",
        "    text_3bw =text_3\n",
        "\n",
        "    descriptor=text_0\n",
        "    return [text_0,text_0bw,text_1,text_1bw,text_2,text_2bw,text_3,text_3bw],descriptor"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "45b-6_CguTMA"
      },
      "source": [
        "## small mesh3 example"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YE0e_rAfxwbq"
      },
      "outputs": [],
      "source": [
        "print (f2.shape)\n",
        "v2= orig_v_template.unsqueeze(0)#Tshape\n",
        "#v2=output.vertices.detach()#.cpu()#.numpy().squeeze()\n",
        "print (v2.device,v2.shape,v2.dtype,f2.device,f2.shape,f2.dtype)\n",
        "print (f2.max(),f2.min(),f2[0,0])\n",
        "\n",
        "with torch.no_grad():\n",
        "  mesh3 = Meshes(v2*3-0.3,f2)\n",
        "  mesh3.textures = TexturesVertex(verts_features=v2*0) #\n",
        "  #mesh3.textures =texture#\n",
        "\n",
        "  l_e_orig = mesh_edge_loss(mesh3) \n",
        "  l_l_orig =  mesh_laplacian_smoothing(mesh3)\n",
        "  l_n_orig =  mesh_normal_consistency(mesh3)\n",
        "\n",
        "with torch.no_grad():\n",
        "  lights = canonicL#PointLights(device=device, location=[[0.0, 0.0, -3.0]])#hack to reset to normal light\n",
        "  images = renderer(mesh3, lights=lights)\n",
        "  print (images.shape,images.max(),images.min())#directly 0,1 -- ok\n",
        "  plt.figure(figsize=(20,20))\n",
        "  plt.imshow(images[0,:,:,:3].cpu())\n",
        "  plt.show()\n",
        "\n",
        "#raise Exception"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "njEZKtnAoJMH"
      },
      "source": [
        "## new edge loss\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DL-R3hMIoOHd"
      },
      "outputs": [],
      "source": [
        "def mesh_edge_len(meshes):\n",
        "    if meshes.isempty():\n",
        "        return torch.tensor(\n",
        "            [0.0], dtype=torch.float32, device=meshes.device, requires_grad=True\n",
        "        )\n",
        "\n",
        "    N = len(meshes)\n",
        "    edges_packed = meshes.edges_packed()  # (sum(E_n), 3)\n",
        "    verts_packed = meshes.verts_packed()  # (sum(V_n), 3)\n",
        "\n",
        "    edge_to_mesh_idx = meshes.edges_packed_to_mesh_idx()  # (sum(E_n), )\n",
        "    num_edges_per_mesh = meshes.num_edges_per_mesh()  # N\n",
        "\n",
        "    verts_edges = verts_packed[edges_packed]\n",
        "    v0, v1 = verts_edges.unbind(1)\n",
        "    l=(v0 - v1).norm(dim=1, p=2) \n",
        "    return l\n",
        "with torch.no_grad():\n",
        "  edge_len = mesh_edge_len(mesh3)#initial mesh\n",
        "print (\"initial edges\",edge_len.shape)\n",
        "\n",
        "wel = edge_len*0+1\n",
        "plt.scatter(edge_len.cpu(),wel.cpu())\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tt6zxOmD-OVD"
      },
      "source": [
        "## multi view defines "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OOrsJUQ5RoBa"
      },
      "outputs": [],
      "source": [
        "usePerspective=False\n",
        "def getCamP(R,T):\n",
        "    if usePerspective:\n",
        "        return FoVPerspectiveCameras(device=device, R=R, T=T,fov=30)##fov like zoom? calibrate negates\n",
        "\n",
        "    return FoVOrthographicCameras(device=device, R=R, T=T)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xEQ9_m-nAAOS"
      },
      "outputs": [],
      "source": [
        "fv=50+40#larger fov and close to camera -- large distort\n",
        "cdist=2.+1#if small -- more distort\n",
        "\n",
        "num_views = 360#used in target cameras only\n",
        "\n",
        "easc=15\n",
        "elev = torch.linspace(-easc,+easc, num_views//2)\n",
        "elev=torch.cat([elev,torch.flip(elev,[0])])\n",
        "\n",
        "TaT = 0#hmm, only visible when having different angle than sun?\n",
        "asc=0.001#no camera change gkobally - mesh will be rotated instead\n",
        "azim = torch.linspace(TaT-asc,TaT+asc, num_views//2)\n",
        "if False:\n",
        "    azim=torch.cat([azim,torch.flip(azim,[0])])\n",
        "else:\n",
        "    azim=torch.cat([azim,azim])\n",
        "R2, T2 = look_at_view_transform(dist=cdist, elev=elev, azim=azim)\n",
        "\n",
        "#cameras =FoVPerspectiveCameras(device=device, R=R2, T=T2,fov=fv)##pespective? orthotraphic finer control with perspective, e.g. translation in ortho weird\n",
        "num_views = train_num_views#4#8\n",
        "   \n",
        "#random  views\n",
        "def sampleRC(num_views=num_views):\n",
        "  if True:#sample in range\n",
        "    elev = torch.zeros(num_views).uniform_(-easc,+easc)\n",
        "    #asc=180#overwrite---\n",
        "    azim = torch.zeros(num_views).uniform_(TaT-asc,TaT+asc)#try narrow angle\n",
        "  else:\n",
        "    elev = torch.linspace(16,17, num_views)\n",
        "    azim = torch.linspace(-360-40,-360+40, num_views)\n",
        "  R2, T2 = look_at_view_transform(dist=cdist, elev=elev, azim=azim)\n",
        "  cameras =getCamP(R=R2, T=T2)\n",
        "  return cameras"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M7IvK0fpizbW"
      },
      "source": [
        "# ML Stuff"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zV5JuGZYxrck"
      },
      "source": [
        "## Sobolev preconditioner "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Cx9ayl9YRoBc"
      },
      "outputs": [],
      "source": [
        "def norm_laplacian2(verts_packed,edges_packed):\n",
        "    L= norm_laplacian(verts_packed,edges_packed)\n",
        "    V=L.shape[0]\n",
        "    # Add the diagonal indices\n",
        "    vals = torch.sparse.sum(L, dim=0).to_dense()\n",
        "    indices = torch.arange(V, device='cuda')\n",
        "    idx = torch.stack([indices, indices], dim=0)\n",
        "    L = torch.sparse.FloatTensor(idx, vals, (V, V)) - L\n",
        "    return L"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EuKDvbATxu2B"
      },
      "outputs": [],
      "source": [
        "verts_packed = orig_v_template.squeeze()#mesh3.verts_packed()  # (sum(V_n), 3)\n",
        "edges_packed = mesh3.edges_packed()  \n",
        "from largesteps.solvers import ConjugateGradientSolver,solve,CholeskySolver\n",
        "if bSOBOLEV:    \n",
        "    #from solvers import *\n",
        "    #I0 = torch.ones(verts_packed.shape[0])\n",
        "    #I = torch.diag(I0).to_sparse().cuda()\n",
        "    V=verts_packed.shape[0]\n",
        "    idx=torch.arange(V).cuda()\n",
        "    idx = torch.stack([idx,idx], dim=0)\n",
        "    I = torch.sparse.FloatTensor(idx, torch.ones(V).cuda(), (V, V))\n",
        "    if not bWeightedEdges:#normed with degree\n",
        "        L = -laplacian(verts_packed,edges_packed)#so diagonal is 1\n",
        "    else: #inv distance - so small polys smoother\n",
        "        L= 1e-2*norm_laplacian2(verts_packed,edges_packed)\n",
        "    M=I+fSobolevStrength*L#\n",
        "    solver=ConjugateGradientSolver(M.coalesce())#\n",
        "    #solver = CholeskySolver(M.coalesce())\n",
        "    def applyPreconditioner(x):\n",
        "      return solve(solver,x)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VWbgXjrRkK4Y"
      },
      "outputs": [],
      "source": [
        "def saveI(img,i):\n",
        "  img=np.uint8(img*255)\n",
        "  imageio.imwrite(\"frames2/\"+str(i) + '.jpg',img,quality=95)\n",
        "\n",
        "def saveImageSet(mesh3,images=[],text_to_add= \"just a T\"):\n",
        "  start_offset= len(os.listdir(\"frames2/\"))#in case adding to other frames\n",
        "  print (\"start offset\",start_offset)\n",
        "  images=torch.cat(images)\n",
        "  for i in range(images.shape[0]):\n",
        "    saveI(images[i,:,:,:3].cpu().numpy(),i+start_offset)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9XDDRC5cwzHE"
      },
      "source": [
        "## CLIP"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lU-CoFssw28N"
      },
      "outputs": [],
      "source": [
        "#Import CLIP and load the model\n",
        "from CLIP import clip\n",
        "print (clip.available_models())\n",
        "clip_preprocess=  torchvision.transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))\n",
        "['RN50', 'RN101', 'RN50x4', 'RN50x16', 'RN50x64', 'ViT-B/32', 'ViT-B/16', 'ViT-L/14']\n",
        "if False:#standard openai\n",
        "  perceptor, preprocess = clip.load('ViT-L/14', jit=False)#32 initially\n",
        "  perceptor.eval().requires_grad_(False);\n",
        "  clip.available_models()\n",
        "elif True:#combine 2 models\n",
        "  perceptor0,_ = clip.load('ViT-B/16', jit=False)#32 initially\n",
        "  perceptor1,_ = clip.load('ViT-B/32', jit=False)#any difference?##arghh, uses other size!!! -- 288!\n",
        "  perceptor0.eval().requires_grad_(False)\n",
        "  perceptor1.eval().requires_grad_(False)\n",
        "    \n",
        "  print (perceptor0.visual.input_resolution )\n",
        "  print (perceptor1.visual.input_resolution )\n",
        " \n",
        "  class CompoClip(object):\n",
        "    def __init__(self):\n",
        "        self.p0=perceptor0\n",
        "        self.p1=perceptor1\n",
        "        \n",
        "    def encode_text(self,t):\n",
        "        return torch.cat([self.p0.encode_text(t),self.p1.encode_text(t)],-1)\n",
        "          \n",
        "    def encode_image(self,t):\n",
        "        return torch.cat([self.p0.encode_image(t),self.p1.encode_image(t)],-1)\n",
        "          \n",
        "  perceptor=CompoClip()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2mgAA8jQx2mG"
      },
      "outputs": [],
      "source": [
        "pbatch_size=1\n",
        "\n",
        "#replace the B H W C=4 channels and pass to CLIP\n",
        "def procClip(x):\n",
        "  return x.permute(0,3,1,2)[:,:3]\n",
        "\n",
        "def soo(sideX):\n",
        "  size = int((      0.7+np.random.rand()*0.3  )*sideX)\n",
        "  offsetx = torch.randint(0, int(sideX - size), ())\n",
        "  offsety = torch.randint(0, int(sideX - size), ())\n",
        "  return offsetx,offsety,size \n",
        "\n",
        "#gives crops patches proportional to size\n",
        "def getP(cutn,sideX):\n",
        "  p_s = []#just patch crops\n",
        "  for ch in range(cutn):\n",
        "      offsetx,offsety,size =soo(sideX)\n",
        "      p_s.append((offsetx,offsety,size))\n",
        "  return p_s\n",
        "\n",
        "#can work with 1 image only\n",
        "def patch(into,cutn=32):\n",
        "  #up_noise=0.1\n",
        "  scaler=1\n",
        "  patch = getP(cutn,into.shape[2])\n",
        "  p_s=[]\n",
        "  for ch in range(cutn):\n",
        "      offsetx,offsety,size = patch[ch]\n",
        "      apper = into[:, :, offsetx:offsetx + size, offsety:offsety + size]\n",
        "      apper = torch.nn.functional.interpolate(apper, (int(224*scaler), int(224*scaler)), mode='bilinear', align_corners=True)\n",
        "      p_s.append(apper)\n",
        "\n",
        "  into = torch.cat(p_s, 0)\n",
        "  return into\n",
        "\n",
        "#loss for text embedding t,t2 into image into; orig gives optionally mask where all 1s in three channels means to ignore!\n",
        "def lossClip(into,t,t2=None,orig=None,use_patch=False):\n",
        "  if orig is None:\n",
        "    orig =into\n",
        "  if use_patch or into.shape[2]!=224:\n",
        "    into1=patch(into,cutn=extra_patch_clip)\n",
        "    into=into1\n",
        "  else:\n",
        "    into =  torch.nn.functional.interpolate(into, (int(224), int(224)), mode='bilinear', align_corners=True)\n",
        "  if np.random.rand()<0.01:# debug sanity\n",
        "      print (\"clip patch cuts\",into.shape,\"texts\",len(t))\n",
        "    \n",
        "  into = clip_preprocess(into)\n",
        "  if clip_flip:\n",
        "    into = torch.cat([into,torch.flip(into,dims=[3])])  \n",
        "  #no vert into = torch.cat([into,torch.flip(into,dims=[2])])  \n",
        "  #print (\"inti\",into.shape)\n",
        "  iii = perceptor.encode_image(into)\n",
        "  l1=  0\n",
        "  for te in t:\n",
        "    v=torch.cosine_similarity(te, iii, -1)\n",
        "    #v=v.mean()\n",
        "    v=0.5*(v.max()+v.mean())#new experiment, only best view activates!!\n",
        "    l1 += -10*v/len(t)\n",
        "  if t2 is not None:\n",
        "    for part in t2:\n",
        "      l1= l1+  1*torch.cosine_similarity(part, iii, -1).mean()/len(t2)\n",
        "  return l1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eTyzlG37IOxz"
      },
      "source": [
        "## symmetry"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MH6xpRz7L1uE"
      },
      "outputs": [],
      "source": [
        "import pickle\n",
        "if useSymmetry:  #try symmetry of mesh\n",
        "  try:\n",
        "    ixs=pickle.load(open('symVerts%d.dat'%(orig_v_template.shape[0]),'rb'))\n",
        "    assert(len(ixs)==orig_v_template.shape[0])\n",
        "  except:\n",
        "      mir = orig_v_template*1\n",
        "      mir[:,0]*=-1#only x coordinate!\n",
        "      print (mir.shape)\n",
        "      out=[]\n",
        "      ixs=[]\n",
        "      for i in range(orig_v_template.shape[0]):\n",
        "        d = orig_v_template[i:i+1]-mir\n",
        "        dd= (d**2).sum(1)\n",
        "        m=dd.min()\n",
        "        out.append(m.item())\n",
        "        ixs.append(np.argmin(dd.cpu()))#closest mirror to vertex i\n",
        "      out=np.array(out)\n",
        "      pickle.dump(ixs,open('symVerts%d.dat'%(orig_v_template.shape[0]),'wb'))\n",
        "      print (\"SYM verts\",out.mean(),out.min(),out.max())\n",
        "  def sym(v):\n",
        "    v2 = v[ixs,:]#for each vertex gives position of its mirror\n",
        "    m = torch.cat([v2[:,:1]*-1,v2[:,1:]],1)#should have grad, project mirror across x axis\n",
        "    return (m+v)/2\n",
        "\n",
        "  test=sym(orig_v_template)\n",
        "else:\n",
        "  def sym(v):\n",
        "    return v#dummy"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KeWlyoXCwzb5"
      },
      "source": [
        "## spectral features\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7dFemdUSkx8n",
        "scrolled": false
      },
      "outputs": [],
      "source": [
        "#emebedding of mesh - -takes 1 minute, disable if not using\n",
        "LAP_EMB = True\n",
        "blockGC = False#if  True simple bias stuff!!\n",
        "NE=120#50#try 50 or 150\n",
        "bUseClust = False\n",
        "loaded=False\n",
        "\n",
        "if LAP_EMB:\n",
        "  try:\n",
        "      lname='meshEmb_L%d_%d.pt'%(HRL,NE+int(bUseClust)*22)\n",
        "      plname='%s/'%(resourcePath) + lname\n",
        "      mesh_embed= torch.load(plname)\n",
        "      print (\"loaded\",mesh_embed.shape)\n",
        "      mesh_embed=mesh_embed.cuda()\n",
        "      \n",
        "      NE = mesh_embed.shape[1]\n",
        "      assert (mesh_embed.shape[0]==orig_v_template.shape[0])\n",
        "      print (\"loading worked\")\n",
        "      loaded=True\n",
        "      #assert (mesh_embed.shape[1]==NE)\n",
        "  except Exception as e:\n",
        "      print (\"could not load\",e,orig_v_template.shape)\n",
        "      print (e)\n",
        "      #recalc\n",
        "      #raise Exception\n",
        "else:\n",
        "  mesh_embed=1*orig_v_template\n",
        "  NE =3\n",
        "\n",
        "if LAP_EMB and not loaded:\n",
        "  useSelf=True\n",
        "  print (\"use just one mesh clone if identical, no need for full embedding\")\n",
        "  f_02 = f2[:,:f2.shape[1]//NClones]\n",
        "  s=f_02.max().item()+1\n",
        "  print (f2.shape,f_02.shape,s,orig_v_template.shape)\n",
        "  from scipy.sparse import lil_matrix\n",
        "  A = lil_matrix((s,s))\n",
        "  print (\"A array\",A.shape,A.dtype)\n",
        "\n",
        "  vals=[]\n",
        "  for i in range(f_02.shape[1]):\n",
        "                  for j in range(3):\n",
        "                      for jj in range(3):\n",
        "                          if j !=jj:#self similarity or not\n",
        "                              vals.append(1)\n",
        "                              A[f_02[0,i,j].cpu(),f_02[0,i,jj].cpu()]=vals[-1]\n",
        "                      if useSelf:\n",
        "                          A[f_02[0,i,j].cpu(),f_02[0,i,j].cpu()]=0.0001 \n",
        "  plt.plot(np.sort(np.array(vals[::10])))\n",
        "  plt.title(\"adjacency vals\")\n",
        "  plt.show()\n",
        "\n",
        "  from sklearn.manifold import SpectralEmbedding\n",
        "  tsne = SpectralEmbedding(NE//NClones,affinity ='precomputed')\n",
        "  assert(NE%NClones ==0)\n",
        "  NE0=NE//NClones\n",
        "  zx=torch.FloatTensor(tsne.fit_transform(A)).cuda()#\n",
        "  mesh_embed= torch.zeros(zx.shape[0]*NClones,NE).cuda()#zx\n",
        "  for i in range(NClones):#copy blocks\n",
        "    mesh_embed[i*zx.shape[0]:(i+1)*zx.shape[0],i*NE0:(i+1)*NE0]=zx\n",
        "  del A\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EupFlZByRoBl"
      },
      "outputs": [],
      "source": [
        "if LAP_EMB and not loaded:\n",
        "    torch.save(mesh_embed.cpu(),plname)\n",
        "    \n",
        "NE=NE+3\n",
        "mesh_embed= torch.cat([mesh_embed,orig_v_template],1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Lj40nKSccqhC"
      },
      "outputs": [],
      "source": [
        "mesh_embed-=mesh_embed.mean(0).unsqueeze(0)\n",
        "mesh_embed /= mesh_embed.std(0).unsqueeze(0)\n",
        "plt.plot(mesh_embed[:,:9].cpu())\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LWjrw9N9bhEo"
      },
      "outputs": [],
      "source": [
        "if LAP_EMB:\n",
        "  plt.figure(figsize=(27,27))\n",
        "  for z in range(9):\n",
        "    plt.subplot(3,3,z+1)\n",
        "    plt.scatter(orig_v_template.cpu()[::4,0],orig_v_template.cpu()[::4,1],c=mesh_embed[::4,z].cpu(),s=3,alpha=0.4)\n",
        "    plt.axis('off')\n",
        "  plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Graph Conv"
      ],
      "metadata": {
        "id": "kU85o-S2Du0K"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from pytorch3d.ops.graph_conv import GraphConv,gather_scatter\n",
        "\n",
        "class GraphConv(nn.Module):\n",
        "    \"\"\"A single graph convolution layer. with symmtric normalize\"\"\"\n",
        "    def __init__(\n",
        "        self,\n",
        "        input_dim: int,\n",
        "        output_dim: int,\n",
        "        init: str = \"normal\",\n",
        "        directed: bool = False,\n",
        "    ) -> None:\n",
        "        super().__init__()\n",
        "        self.input_dim = input_dim\n",
        "        self.output_dim = output_dim\n",
        "        self.directed = directed\n",
        "        self.w0 = nn.Linear(input_dim, output_dim)\n",
        "        self.w1 = nn.Linear(input_dim, output_dim)\n",
        "\n",
        "        if init == \"normal\":\n",
        "            nn.init.normal_(self.w0.weight, mean=0, std=0.01)\n",
        "            nn.init.normal_(self.w1.weight, mean=0, std=0.01)\n",
        "            # pyre-fixme[16]: Optional type has no attribute `data`.\n",
        "            self.w0.bias.data.zero_()\n",
        "            self.w1.bias.data.zero_()\n",
        "        elif init == \"zero\":\n",
        "            self.w0.weight.data.zero_()\n",
        "            self.w1.weight.data.zero_()\n",
        "        else:\n",
        "            raise ValueError('Invalid GraphConv initialization \"%s\"' % init)\n",
        "\n",
        "    def forward(self, verts, edges,deg):\n",
        "        if verts.is_cuda != edges.is_cuda:\n",
        "            raise ValueError(\"verts and edges tensors must be on the same device.\")\n",
        "        if verts.shape[0] == 0:\n",
        "            # empty graph.\n",
        "            return verts.new_zeros((0, self.output_dim)) * verts.sum()\n",
        "\n",
        "        verts_w0 = self.w0(verts)  # (V, output_dim)\n",
        "        verts_w1 = self.w1(verts*deg)  # (V, output_dim)\n",
        "\n",
        "        if torch.cuda.is_available() and verts.is_cuda and edges.is_cuda:\n",
        "            neighbor_sums = gather_scatter(verts_w1, edges, self.directed)\n",
        "        else:\n",
        "            neighbor_sums = gather_scatter_python(\n",
        "                verts_w1, edges, self.directed\n",
        "            )  # (V, output_dim)\n",
        "\n",
        "        # Add neighbor features to each vertex's features.\n",
        "        out = verts_w0 + deg*neighbor_sums\n",
        "        return out\n",
        "    \n",
        "def deg(V,edges):\n",
        "    e0, e1 = edges.unbind(1)\n",
        "    idx01 = torch.stack([e0, e1], dim=1)  # (E, 2)\n",
        "    idx10 = torch.stack([e1, e0], dim=1)  # (E, 2)\n",
        "    idx = torch.cat([idx01, idx10], dim=0).t()  # (2, 2*E)\n",
        "    ones = torch.ones(idx.shape[1], dtype=torch.float32).cuda()\n",
        "    A = torch.sparse.FloatTensor(idx, ones, (V, V))\n",
        "    # the sum of i-th row of A gives the degree of the i-th vertex\n",
        "    deg = torch.sparse.sum(A, dim=1).to_dense()\n",
        "    return deg\n",
        "\n",
        "#multiply by diagonal matrix\n",
        "def ngc(s,g,x):\n",
        "          return g(x,s.edges,s.deg)"
      ],
      "metadata": {
        "id": "MKyHdJXiDwuu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LYPgKJhfre8D"
      },
      "source": [
        "## MLP mesh deform"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oIcLfPThRoBo"
      },
      "outputs": [],
      "source": [
        "class PlainDeform(torch.nn.Module):\n",
        "  def __init__(self,orig_v,extradim=0,gcdepth=None,nZ=50*3,useClip=0,initGamma=5e-2,initPV=1):\n",
        "    super().__init__()\n",
        "    #hack: *4 when using plain initially\n",
        "    self.w = nn.Parameter(initPV*(orig_v*0).uniform_(-0.04,0.04))    #pure bias\n",
        "    \n",
        "    self.w1 = nn.Parameter(torch.zeros(NE,3).uniform_(-0.005,0.005))#linear on spectral features#this is too wild\n",
        "    \n",
        "    self.mlp = nn.Sequential(nn.Linear(NE,nZ),nn.ReLU(True),nn.Linear(nZ,nZ),nn.ReLU(True),nn.Linear(nZ,3))\n",
        "    self.gamma0=nn.Parameter(torch.zeros(1)+1)\n",
        "    self.gamma=nn.Parameter(torch.zeros(1)+initGamma)\n",
        "    self.bSYM=useSymmetry\n",
        "\n",
        "    if graphConv:\n",
        "        z=32*1\n",
        "        cin=3+NE\n",
        "        self.nlin=nn.LeakyReLU(0.2)\n",
        "        self.first=GraphConv(cin,z)\n",
        "        self.edges = mesh3.edges_packed()\n",
        "        self.deg=deg(orig_v_template.shape[0],self.edges).unsqueeze(1)\n",
        "        self.deg = 1/torch.sqrt(self.deg)\n",
        "        print (\"edges\",self.edges.shape,self.deg.shape)\n",
        "        self.gc = nn.ModuleList()\n",
        "        for i in range(nGraphLayers):\n",
        "              self.gc.append(GraphConv(z,z))#was +6\n",
        "        self.gc.append(GraphConv(z,3))\n",
        "    \n",
        "  def forward(self,v_posed):\n",
        "    out=self.gamma0*self.w+ self.gamma*self.mlp(mesh_embed)+self.gamma*mesh_embed@self.w1    \n",
        "    if graphConv:\n",
        "        if self.bSYM:\n",
        "          out=sym(out)\n",
        "        z=ngc(self,self.first,torch.cat([mesh_embed,out],1))    \n",
        "        br=0\n",
        "        for g in self.gc:\n",
        "          oldz=z\n",
        "\n",
        "          z=self.nlin(z)\n",
        "          if True and br !=len(self.gc)-1:\n",
        "              z=oldz+fgraphnet*ngc(self,g,z)\n",
        "          else:\n",
        "              z=ngc(self,g,z)\n",
        "          br +=1\n",
        "        out=out+fgraphnet*z \n",
        "\n",
        "    if bSOBOLEV:#u reparametrize#ahaa try also full code, with the v_shaped; this is the from_differential command in the EPFL paper TODO solved of EPFL\n",
        "      out = applyPreconditioner(out)\n",
        "    if self.bSYM:\n",
        "      out=sym(out)##          \n",
        "    out=out+v_posed[:,:3]\n",
        "    self.last=out-v_posed[:,:3]#orig_v_template#\n",
        "    return out\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SimWWyuVuuhz"
      },
      "source": [
        "## helper network archis"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CbRSh9pZum2r"
      },
      "outputs": [],
      "source": [
        "def double_conv(in_channels, out_channels):\n",
        "    return nn.Sequential(\n",
        "        nn.Conv2d(in_channels, out_channels, 3, padding=1),\n",
        "        nn.ReLU(inplace=True),\n",
        "        nn.InstanceNorm2d(out_channels),\n",
        "        nn.Conv2d(out_channels, out_channels, 3, padding=1),\n",
        "        nn.ReLU(inplace=True)\n",
        "    )   \n",
        "\n",
        "class UNet(nn.Module):\n",
        "    def __init__(self, n_class,c_in=3,c=88):\n",
        "        super().__init__()\n",
        "        self.c=c\n",
        "\n",
        "        self.dconv_down1 = double_conv(c_in, c)\n",
        "        self.dconv_down2 = double_conv(c, 2*c)\n",
        "        self.dconv_down3 = double_conv(2*c, 4*c)\n",
        "        self.dconv_down4 = double_conv(4*c, 8*c)        \n",
        "\n",
        "        self.maxpool = lambda x:F.interpolate(x, scale_factor=0.5, mode='bilinear')#nn.AvgPool2d(2)#nn.MaxPool2d(2)\n",
        "        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)        \n",
        "        \n",
        "        self.dconv_up3 = double_conv(4*c + 8*c, 4*c)\n",
        "        self.dconv_up2 = double_conv(2*c + 4*c, 2*c)\n",
        "        self.dconv_up1 = double_conv(c + 2*c, c)\n",
        "        self.conv_last = nn.Conv2d(c+c_in, 3, 1)\n",
        "\n",
        "    #z is noise -- what scale?\n",
        "    def forward(self, x_,zraw=None):\n",
        "        conv1 = self.dconv_down1(x_)\n",
        "        x = self.maxpool(conv1)\n",
        "\n",
        "        conv2 = self.dconv_down2(x)\n",
        "        x = self.maxpool(conv2)\n",
        "        \n",
        "        conv3 = self.dconv_down3(x)\n",
        "        x = self.maxpool(conv3)   \n",
        "        \n",
        "        x = self.dconv_down4(x)\n",
        "        \n",
        "        x = self.upsample(x)        \n",
        "        x = torch.cat([x, conv3], dim=1)\n",
        "        \n",
        "        x = self.dconv_up3(x)\n",
        "        x = self.upsample(x)        \n",
        "        x = torch.cat([x, conv2], dim=1)       \n",
        "\n",
        "        x = self.dconv_up2(x)\n",
        "        x = self.upsample(x)        \n",
        "        x = torch.cat([x, conv1], dim=1)   \n",
        "        \n",
        "        x = self.dconv_up1(x)\n",
        "        out = self.conv_last(torch.cat([x,x_],1))#new Dec. 2 -- also add last\n",
        "        return out"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tVUM59RyFFt2"
      },
      "source": [
        "# classes and functions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aUuvPR0FIjg9"
      },
      "source": [
        "## Class for texture learning"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "INMALW3wRoBq"
      },
      "outputs": [],
      "source": [
        "#full grid structure\n",
        "TEXSIZEw=TEXSIZE#*NClones\n",
        "def getImageL():\n",
        "  ix=[]\n",
        "  v=[]\n",
        "  V=TEXSIZE*TEXSIZEw\n",
        "  sk=0\n",
        "  for h in range(TEXSIZE-1):\n",
        "    for w in range(TEXSIZEw-1):\n",
        "      \n",
        "      i=h*TEXSIZEw + w#curr index\n",
        "\n",
        "      j=i+1#hor neighbour\n",
        "      \n",
        "      ix.append([i,j])\n",
        "      v.append(-1.0)\n",
        "      ix.append([j,i])\n",
        "      v.append(-1.0)\n",
        "      ix.append([j,j])\n",
        "      v.append(1.0)\n",
        "      ix.append([i,i])\n",
        "      v.append(1.0)\n",
        "\n",
        "      j=i+TEXSIZEw#ver neighbour\n",
        "\n",
        "      ix.append([i,j])\n",
        "      v.append(-1.0)\n",
        "      ix.append([j,i])\n",
        "      v.append(-1.0)\n",
        "      ix.append([j,j])\n",
        "      v.append(1.0)\n",
        "      ix.append([i,i])\n",
        "      v.append(1.0)\n",
        "  \n",
        "  ix= np.array(ix).T#so 2 x N\n",
        "  print (\"skipped\",sk,\"did indices in image laplacian\",ix.shape)##non-full -- still solve?\n",
        "  #8x entries added per valid pixel; symmetry and back - so 4; diagonals i,j and ij,ji\n",
        "  v=np.array(v)\n",
        "  L= torch.sparse_coo_tensor(ix, v, (V,V)).coalesce().cuda()\n",
        "\n",
        "  ix=np.arange(V)\n",
        "  ix=np.int32(np.array([ix,ix]))#so 2xN\n",
        "  v=np.ones(V)\n",
        "  print (ix.shape,v.shape)\n",
        "  I= torch.sparse_coo_tensor(ix, v, (V,V)).coalesce().cuda()\n",
        "  return L,I\n",
        "solveTex=True\n",
        "if solveTex:\n",
        "  try:\n",
        "      IL,II = torch.load('textureLaplacian%d.dat'%(TEXSIZEw))\n",
        "      print (\"loaded texture Laplacian\")\n",
        "  except:\n",
        "      print (\"calculating texture laplacian\")\n",
        "      IL,II=getImageL()\n",
        "      torch.save((IL,II ),'textureLaplacian%d.dat'%(TEXSIZEw))\n",
        "  print (IL.shape,II.shape)\n",
        "  IM=II+(.8+0)*IL\n",
        "  solverI = CholeskySolver(IM.coalesce().float())\n",
        "\n",
        "  IM=II+(.8)*IL\n",
        "  solverIN = CholeskySolver(IM.coalesce().float())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7YlxPV7OIlQo"
      },
      "outputs": [],
      "source": [
        "class TextureL(nn.Module):\n",
        "    def __init__(self,c=6):\n",
        "      super().__init__()\n",
        "      tex=torch.zeros(1,TEXSIZE,TEXSIZEw,3).uniform_(-3,3)   # .uniform_(0.25,0.75)\n",
        "      self.tex = nn.Parameter(tex)\n",
        "      texn=torch.zeros(1,TEXSIZE,TEXSIZEw,3).uniform_(-0.2,0.2)\n",
        "      self.texn = nn.Parameter(texn)\n",
        "\n",
        "      if True:\n",
        "          self.unet0= nn.Conv2d(3,3,7,1,3)\n",
        "          self.unet1= nn.Conv2d(3,3,7,1,3)\n",
        "     \n",
        "    def forward(self,x=None,z=None):\n",
        "      tex=torch.sigmoid(self.tex)\n",
        "      texn=self.texn\n",
        "        \n",
        "      if solveTex:\n",
        "        tex_ = solve(solverI,tex.contiguous().view(-1,3))\n",
        "        tex=tex_.view(tex.shape)\n",
        "        texn_ = solve(solverIN,texn.contiguous().view(-1,3))\n",
        "        texn=texn_.view(texn.shape)\n",
        "\n",
        "        if True:\n",
        "          tex=tex+0.15*self.unet0(tex.permute(0,3,1,2)).permute(0,2,3,1)\n",
        "          texn=texn+0.15*self.unet1(texn.permute(0,3,1,2)).permute(0,2,3,1)\n",
        "        \n",
        "      return tex,texn\n",
        "     "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5YJ0K4Ak-6vr"
      },
      "source": [
        "## renderer functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GJ4PPIEK0Az5"
      },
      "outputs": [],
      "source": [
        "from pytorch3d.renderer.utils import TensorProperties\n",
        "from typing import Optional,Union\n",
        "from pytorch3d.renderer.mesh.rasterizer import Fragments\n",
        "from pytorch3d.ops import interpolate_face_attributes\n",
        "from pytorch3d.renderer.mesh.shading import _apply_lighting\n",
        "from pytorch3d.renderer.blending import softmax_rgb_blend\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### shadows"
      ],
      "metadata": {
        "id": "Uf6gYm_7-jGF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "R_, T_ = look_at_view_transform(dist=5, elev=45, azim=0)\n",
        "#if 0 - 0,0,10 light\n",
        "#if 45 -> 0,1,1 light\n",
        "canCam =getCamP(R=R_, T=T_)\n",
        "class ShaderWithDepth(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()        \n",
        "    def forward(self, meshes_world,size=None,**kwargs) -> torch.Tensor:\n",
        "        raster_settings_shadow = RasterizationSettings(\n",
        "            image_size=size, \n",
        "            blur_radius=np.log(1. / 1e-4 - 1.)*1e-5,\n",
        "            faces_per_pixel=1, perspective_correct=usePerspective and False,max_faces_per_bin=orig_v_template.shape[0])#\n",
        "        rasterizer = MeshRasterizer(cameras=canCam, raster_settings=raster_settings_shadow)     \n",
        "        fragments = rasterizer(meshes_world)##TODO pass raster_settings to overwrite\n",
        "        return fragments.zbuf#TODO any way to smooth?\n",
        "    \n",
        "depthShader = ShaderWithDepth()"
      ],
      "metadata": {
        "id": "kOADlz55-iGI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "NP=25\n",
        "a=[]\n",
        "for i in [-2,-1,0,1,2]:\n",
        "  for j in [-2,-1,0,1,2]:\n",
        "    a.append(torch.FloatTensor([i,j]).unsqueeze(0))\n",
        "off = torch.cat(a)\n",
        "print (off)\n",
        "off = off.cuda().view(NP,1,1,1,2)/sideX\n",
        "#off = torch.zeros(NP,1,1,1,2)\n",
        "#off=off.cuda().uniform_(-1,1)/500#poisson disc\n",
        "\n",
        "def calcShadows(meshes,texels,fragments):\n",
        "     verts = meshes.verts_packed()  # (V, 3)\n",
        "     faces = meshes.faces_packed()  # (F, 3)\n",
        "    \n",
        "     adjust = 1.5/(verts[:,1].max()-verts[:,1].min())#so how streched\n",
        "     with torch.no_grad():\n",
        "            depth = (depthShader(meshes,size=texels.shape[1])-1)/99.0\n",
        "            #plt.figure(figsize=(14,14))\n",
        "            #plt.imshow(depth[0,:,:,0].cpu())\n",
        "            #plt.show()\n",
        "            xyd = canCam.transform_points(verts)\n",
        "            xy=xyd[...,:2]*-1#-1 1 ok, but any flipping required?\n",
        "            d=xyd[...,2:3].view(-1)#depth per vertex, of mesh when seen from light, including occluded backsides\n",
        "            shadow_coords =0\n",
        "            for p in range(NP):\n",
        "                zbuf= F.grid_sample(depth.permute(0,3,1,2),xy.view(depth.shape[0],-1,1,2)+off[p])#shadow map values at location\n",
        "                #print (\"zbuf\",zbuf.shape,\"d\",d.shape)#zbuf torch.Size([2, 1, 41853, 1]) d torch.Size([83706])\n",
        "                zbuf = zbuf.view(-1)\n",
        "                delta=zbuf-d#so 0 when exactly lit, when negative . some occlusio\n",
        "                delta*=adjust#so standard scale\n",
        "                delta[delta >-6e-4]=0#where slightly off\n",
        "                fdelta=delta.view(-1,1)[faces]\n",
        "                \n",
        "                shadow_coords_ = interpolate_face_attributes(fragments.pix_to_face, fragments.bary_coords, fdelta)\n",
        "                shadow_coords_ =  torch.exp(3e3*(shadow_coords_))#smoother in low res?   \n",
        "\n",
        "                shadow_coords +=shadow_coords_/NP \n",
        "                #shadow_coords=torch.clamp(shadow_coords,0,1)\n",
        "                \n",
        "            if np.random.rand() <0.1:\n",
        "                print (\"rendered shadow\",shadow_coords.mean())\n",
        "            return shadow_coords"
      ],
      "metadata": {
        "id": "dNreLpbvvhKV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kr6VY7CP1JFK"
      },
      "source": [
        "### new phong shader with new face normals\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QsgKx9HgxDzC"
      },
      "outputs": [],
      "source": [
        "regEye=torch.eye(2).unsqueeze(0).cuda()\n",
        "def getTanSpace(mesh,vertex3,normals):\n",
        "    out=[]\n",
        "    faces_uvs=mesh.textures.faces_uvs_padded()[0].squeeze()##problem -- if other texture vertex textzre... TODO use just plain UVs!\n",
        "    f2=mesh.faces_padded()[0]#so that proper crops of parts work\n",
        "    for n in range(vertex3.shape[0]*0+1):\n",
        "        E1= vertex3[n,f2[:,0]]-vertex3[n,f2[:,1]]\n",
        "        E2= vertex3[n,f2[:,0]]-vertex3[n,f2[:,2]]#1 Nfaces 3     \n",
        "        E12=torch.cat([E1.unsqueeze(1),E2.unsqueeze(1)],1)#concat, 1 Nfaces x 3x2\n",
        "        \n",
        "        dUV1= vt[0,faces_uvs[:,0]]-vt[0,faces_uvs[:,1]]#2d\n",
        "        dUV2= vt[0,faces_uvs[:,0]]-vt[0,faces_uvs[:,2]]\n",
        "        dUV12= torch.cat([dUV1.unsqueeze(1),dUV2.unsqueeze(1)],1)# concat 1 Nfaces x2 x2        \n",
        "        I=torch.inverse(dUV12+1e-10*regEye)\n",
        "        \n",
        "        TB = I.bmm(E12)\n",
        "        TB=F.normalize(TB,dim=-1)#make sure unit vectors!        \n",
        "        N= normals[n]# from normals        \n",
        "        NTB= torch.cat([N.unsqueeze(1),TB],1)#so all together, should be totalfaces x 3x3 -- how to transpose?\n",
        "        out.append(NTB.permute(0,2,1))        \n",
        "    return out#torch.cat(out)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "useShadows"
      ],
      "metadata": {
        "id": "fOFFhsjXdut2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1ry9Vw_h0RCA"
      },
      "outputs": [],
      "source": [
        "#calculates tangent space normals\n",
        "def sample_texturesN(fragments,mesh,texels_nm0):\n",
        "    with torch.no_grad():\n",
        "      NTB=getTanSpace(mesh,mesh.verts_padded(),mesh.faces_normals_padded())    \n",
        "    out=[]\n",
        "    X=texels_nm0##can we calculate first the normals per pixel in some other space? or do in HW space as usual\n",
        "    \n",
        "    #if X is usual texture, here a normal map -- move to tangent space for each\n",
        "    if True:\n",
        "      for n in range(len(T)):#per batch instance\n",
        "          rot=NTB[n*0].view(-1,3,3)#so rotation matrix per face --tangent space        \n",
        "          rotHW = rot[fragments.pix_to_face[n:n+1]]#1HWKx3x3   --fragment for one batch only     \n",
        "          x=X[n:n+1]#for 1 instance , 1KHW3\n",
        "          #print (\"before rotation\",x.shape,rotHW.shape,rot.shape)#before rotation torch.Size([1, 768, 768, 5, 3]) torch.Size([1, 768, 768, 5, 3, 3])\n",
        "          tangent_x=torch.matmul(rotHW.view(-1,3,3),x.contiguous().view(-1,3,1))##will it work; is it correct directopn? does not change size\n",
        "          out.append(tangent_x.view(x.shape))\n",
        "      out=torch.cat(out)\n",
        "      return out#NHWKC\n",
        "    else:#assume all the same -- weird why such errors\n",
        "      rot=NTB[0]\n",
        "      rotHW = rot[fragments.pix_to_face]\n",
        "      #print (\"before rotation\",X.shape,rotHW.shape,rot.shape)\n",
        "      tangent_x=torch.bmm(rotHW.view(-1,3,3),X.contiguous().view(-1,3,1))\n",
        "      return tangent_x.view(X.shape)\n",
        "\n",
        "#combines tangent space nornamls and usual normals\n",
        "#per fragment!\n",
        "def phong_shading_nm(meshes, fragments, lights, cameras, materials, texels,texels_nm) -> torch.Tensor:    \n",
        "    verts = meshes.verts_packed()  # (V, 3)\n",
        "    faces = meshes.faces_packed()  # (F, 3)\n",
        "    vertex_normals = meshes.verts_normals_packed()  # (V, 3)##weird, why not face normals??\n",
        "    if True:#canonical pytorch3d code; use interplation of vertex normals\n",
        "      faces_normals = vertex_normals[faces]\n",
        "      faces_normals=interpolate_face_attributes(fragments.pix_to_face, fragments.bary_coords, faces_normals)\n",
        "    else:\n",
        "      faces_normals = meshes.faces_normals_packed() \n",
        "      faces_normals = faces_normals[fragments.pix_to_face]   \n",
        "    faces_verts = verts[faces]    \n",
        "    pixel_coords = interpolate_face_attributes(fragments.pix_to_face, fragments.bary_coords, faces_verts)    \n",
        "    #normal map as residual to usual pixel normals! -- may need to reshape\n",
        "    pixel_normals = texels_nm+faces_normals\n",
        "\n",
        "    if useShadows and texels.shape[1]>224:#so proper texture, not mono, and also not small partial clip HACK\n",
        "        shadow_coords=calcShadows(meshes,texels,fragments)\n",
        "    else:\n",
        "        shadow_coords=1\n",
        " \n",
        "    ambient, diffuse, specular = _apply_lighting(\n",
        "        pixel_coords, pixel_normals, lights, cameras, materials\n",
        "    )\n",
        "    colors = (ambient + shadow_coords*diffuse) * texels +shadow_coords*specular\n",
        "    \n",
        "    if ORACLE.shaderI is not None:\n",
        "        return colors,F.normalize(pixel_normals,dim=-1)    \n",
        "    return colors,None\n",
        "\n",
        "class SoftPhongShader_nm(nn.Module):\n",
        "    def __init__(\n",
        "        self,\n",
        "        device= \"cpu\",\n",
        "        cameras: Optional[TensorProperties] = None,\n",
        "        lights: Optional[TensorProperties] = None,\n",
        "        materials: Optional[Materials] = None,\n",
        "        blend_params: Optional[BlendParams] = None,\n",
        "    ) -> None:\n",
        "        super().__init__()\n",
        "        self.lights = lights if lights is not None else PointLights(device=device)\n",
        "        self.materials = (\n",
        "            materials if materials is not None else Materials(device=device)\n",
        "        )\n",
        "        self.cameras = cameras\n",
        "        self.blend_params = blend_params if blend_params is not None else BlendParams()\n",
        "\n",
        "    def to(self, device):\n",
        "        # Manually move to device modules which are not subclasses of nn.Module\n",
        "        cameras = self.cameras\n",
        "        if cameras is not None:\n",
        "            self.cameras = cameras.to(device)\n",
        "        self.materials = self.materials.to(device)\n",
        "        self.lights = self.lights.to(device)\n",
        "        return self\n",
        "\n",
        "    def forward(self, fragments: Fragments, meshes: Meshes, **kwargs) -> torch.Tensor:\n",
        "        cameras = kwargs.get(\"cameras\", self.cameras)\n",
        "        if cameras is None:\n",
        "            msg = \"Cameras must be specified either at initialization \\\n",
        "                or in the forward pass of SoftPhongShader\"\n",
        "            raise ValueError(msg)\n",
        "            \n",
        "        try:\n",
        "            mono= meshes.textures.maps_padded()[0,:,:,0].var()<1e-8#hack, so ths is actually BW, ok to set texel constant\n",
        "            useNM=True and useNormalMaps\n",
        "            isUV=True\n",
        "        except:\n",
        "            mono=False\n",
        "            useNM=False#so vertex textures, for mscale hack, and never normal maps\n",
        "            isUV=False\n",
        "        if mono: \n",
        "            if np.random.rand()<10.5:#experiment, no nm for mono \n",
        "              useNM = False#\n",
        "            texels=meshes.textures.maps_padded().view(-1,3).mean(0)#so 3channel rgb\n",
        "            texels=texels.view(1,1,1,1,3)\n",
        "        else:\n",
        "          texels = meshes.sample_textures(fragments)\n",
        "               \n",
        "        materials = kwargs.get(\"materials\", self.materials)\n",
        "        if useNM:\n",
        "            texels_nm0 = meshes.normal_map.sample_textures(fragments)# make sure that the normal map texture is defined!!\n",
        "            texels_nm=sample_texturesN(fragments,meshes,texels_nm0[...,:3])#TODO static method,\n",
        "        else:\n",
        "            texels_nm=0\n",
        "        lights = kwargs.get(\"lights\", self.lights)\n",
        "                \n",
        "        blend_params = kwargs.get(\"blend_params\", self.blend_params)\n",
        "        colors,nrm = phong_shading_nm(\n",
        "            meshes=meshes,\n",
        "            fragments=fragments,\n",
        "            texels=texels,\n",
        "            texels_nm=texels_nm,\n",
        "            lights=lights,\n",
        "            cameras=cameras,\n",
        "            materials=materials,\n",
        "        )#fragment RGB colors, nrm is the optional fragment normalmap normal\n",
        "        znear = kwargs.get(\"znear\", getattr(cameras, \"znear\", 1.0))\n",
        "        zfar = kwargs.get(\"zfar\", getattr(cameras, \"zfar\", 100.0))\n",
        "        images = softmax_rgb_blend(colors, fragments, blend_params, znear=znear, zfar=zfar)\n",
        "        \n",
        "        with torch.no_grad():#avoid color leak beyond 3d figure\n",
        "                mask= 1*images[:,:,:,3:4]\n",
        "                mask[mask>0.1]=1\n",
        "                r=torch.zeros(images.shape[0],1,1,3).uniform_(0,1).cuda()\n",
        "\n",
        "        if fragments.pix_to_face.shape[1]<512:#N, H, W, K\n",
        "            if mono:# or np.random.rand()<0.3:#better textures? cjance to avoud neural shader and pure UV textzre   \n",
        "                if images.shape[2]<400:\n",
        "                    #print (\"rm\",r)\n",
        "                    return torch.cat([mask*images[:,:,:,:3]+(1-mask)*r,images[:,:,:,3:]],3)\n",
        "                return images\n",
        "        \n",
        "        if useNeuralShader and ORACLE.shaderI is not None:\n",
        "            nrmI=softmax_rgb_blend(nrm, fragments, blend_params, znear=znear, zfar=zfar)\n",
        "            try:\n",
        "                sh=0.05*ORACLE.shaderI(torch.cat([images,nrmI],3).permute(0,3,1,2)).permute(0,2,3,1)\n",
        "            except Exception as e:\n",
        "                print (e,\"no neural shader\")\n",
        "                return images\n",
        "          \n",
        "            sh=sh*mask\n",
        "            \n",
        "            ORACLE.shadereg += (sh**2).mean()/mask.mean()\n",
        "            if images.shape[2]<= 400:#add random background\n",
        "              sh =sh+ (1-mask)*r\n",
        "              #print (\"rc\",r)\n",
        "            else:\n",
        "              sh = sh + (1-mask)\n",
        "            return torch.cat([mask*images[:,:,:,:3]+sh,images[:,:,:,3:]],3)\n",
        "        return images"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kIO_rN9MM9AP"
      },
      "source": [
        "### training renderer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lnIJa2CIYb3k"
      },
      "outputs": [],
      "source": [
        "BLP = BlendParams(sigma=1e-4,gamma=1e-4)#will gamma\n",
        "sigma =1e-4*5#*1e-1#*0.3#*4\n",
        "\n",
        "raster_settings_soft = RasterizationSettings(\n",
        "    image_size=224+112*0, \n",
        "    blur_radius=np.log(1. / 1e-4 - 1.)*sigma ,\n",
        "    faces_per_pixel=15, perspective_correct=usePerspective and False,max_faces_per_bin=orig_v_template.shape[0]\n",
        ")##from collab demo -- softer for differentbiable geometry? - -read reference!!\n",
        "\n",
        "class MeshRendererWithDepth(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.rasterizer = MeshRasterizer(cameras=camera, raster_settings=raster_settings_soft )\n",
        "        self.shader = SoftPhongShader_nm(device=device,cameras=camera,lights=lights,blend_params=BLP)\n",
        "        self.fpp = 15-5\n",
        "        \n",
        "    def forward(self, meshes_world,sizeExtra=64*0+96 ,**kwargs) -> torch.Tensor:\n",
        "        #sigma= np.random.rand()*(1e-4-1e-6)+1e-6 #btw 1e-6 and 1e-4\n",
        "        sigma=np.random.rand()*1.5\n",
        "        sigma=10**(-6.+sigma)#so e-6. to e-4.\n",
        "        rs=RasterizationSettings(image_size=224+sizeExtra,blur_radius=np.log(1. / 1e-4 - 1.)*sigma ,\n",
        "        faces_per_pixel=self.fpp, perspective_correct=usePerspective and False,max_faces_per_bin=orig_v_template.shape[0])       \n",
        "        #hmm, is the variable window too bad for memry?\n",
        "        fragments = self.rasterizer(meshes_world,raster_settings=rs, **kwargs)##TODO pass raster_settings to overwrite\n",
        "        images = self.shader(fragments, meshes_world, **kwargs)##is light and camera passed correctly??\n",
        "        return images, fragments.zbuf\n",
        "    \n",
        "class MeshRendererWithDepth2(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.rasterizer = MeshRasterizer(cameras=camera, raster_settings=raster_settings_soft )\n",
        "        self.shader = SoftPhongShader(device=device,cameras=camera,lights=lights,blend_params=BLP)\n",
        "        self.fpp = 28#+18\n",
        "        \n",
        "    def forward(self, meshes_world,sizeExtra=64*0+112 ,**kwargs) -> torch.Tensor:\n",
        "        #sigma= np.random.rand()*(1e-4-1e-6)+1e-6 #btw 1e-6 and 1e-4\n",
        "        sigma=np.random.rand()*1.5\n",
        "        sigma=10**(-6.+sigma)#so e-6. to e-4.\n",
        "        rs=RasterizationSettings(image_size=224+sizeExtra,blur_radius=np.log(1. / 1e-4 - 1.)*sigma ,\n",
        "        faces_per_pixel=self.fpp, perspective_correct=usePerspective and False,max_faces_per_bin=orig_v_template.shape[0])       \n",
        "        #hmm, is the variable window too bad for memry?\n",
        "        fragments = self.rasterizer(meshes_world,raster_settings=rs, **kwargs)##TODO pass raster_settings to overwrite\n",
        "        images = self.shader(fragments, meshes_world, **kwargs)##is light and camera passed correctly??\n",
        "        return images, fragments.zbuf\n",
        "    \n",
        "renderer_train = MeshRendererWithDepth()\n",
        "renderer_train2 = MeshRendererWithDepth2()# no shader or normal stuff, no collision detextion -- less faces!!\n",
        "\n",
        "renderer_train_simple = MeshRenderer(\n",
        "    rasterizer=MeshRasterizer(\n",
        "        cameras=camera, \n",
        "        raster_settings=raster_settings_soft\n",
        "    ),\n",
        "    shader=SoftPhongShader(\n",
        "        device=device, \n",
        "        cameras=camera,\n",
        "        lights=lights\n",
        "    )\n",
        ")#used for clip embeddings inside LSTM deformer\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZJ279d8pwM3h"
      },
      "source": [
        "### inference renderer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OjiAIgs89Fw3",
        "scrolled": false
      },
      "outputs": [],
      "source": [
        "#inference soft too? - difference when makinf videos? what us optimal face count and sigma...\n",
        "\n",
        "raster_settings_soft_highrescor = RasterizationSettings(\n",
        "    image_size=sideX, \n",
        "    blur_radius=np.log(1. / 1e-4 - 1.)*1e-7,\n",
        "    faces_per_pixel=4, perspective_correct=usePerspective,\n",
        "    cull_backfaces=False,\n",
        "    max_faces_per_bin=orig_v_template.shape[0]+10000#for highres\n",
        ")\n",
        "  \n",
        "#higher res, for inference, overwrite other renderer!!\n",
        "renderer = MeshRenderer(\n",
        "    rasterizer=MeshRasterizer(\n",
        "        cameras=camera, \n",
        "        raster_settings=raster_settings_soft_highrescor\n",
        "    ),\n",
        "    shader=SoftPhongShader_nm(\n",
        "        device=device, \n",
        "        cameras=camera,\n",
        "        lights=lights, blend_params=BLP\n",
        "    )\n",
        ")\n",
        "\n",
        "\n",
        "try:\n",
        "    lights=calcLights()\n",
        "    mater=calcMaterials()\n",
        "except Exception as e:\n",
        "    mater=Materials(shininess=100).cuda()\n",
        "    lights=DirectionalLights(ambient_color=((0.2, 0.2, 0.2), ), diffuse_color=((.4, .4, .4), ),\\\n",
        "            specular_color=((.4, .4 ,.4), ), direction=((0, 5, 10), ), device=device)\n",
        "    print (e,\"light error\")\n",
        "  \n",
        "import warnings\n",
        "warnings.filterwarnings(\"default\")   #default (\"error\")\n",
        "\n",
        "tex=torch.zeros(1,TEXSIZE,TEXSIZEw,3).cuda()\n",
        "\n",
        "#mesh3=m_check\n",
        "#mesh3.textures=TexturesUV(maps=tex,faces_uvs=[faces_uvs], verts_uvs=vt)\n",
        "#mesh3.normal_map=mesh3.textures\n",
        "if True:#to debug cameras\n",
        "    with torch.no_grad():       \n",
        "      camera= getCamP(R=R2[:1], T=T2[:1])\n",
        "      r=np.random.randint(R2.shape[0])\n",
        "      camera= getCamP(R=R2[r:r+1], T=T2[r:r+1])\n",
        "      images_ = renderer(mesh3, lights= lights,cameras=camera,materials=mater)\n",
        "      fig=plt.figure(figsize=(10.5,10.5))\n",
        "      plt.imshow(images_[0,:,:,:3].cpu())\n",
        "      plt.show()\n",
        "\n",
        "z=mesh3.verts_packed()[:,2]\n",
        "print (z.shape)\n",
        "plt.plot(np.sort(z.detach().cpu().numpy()))\n",
        "plt.show()\n",
        "\n",
        "images2_ = renderer_train(mesh3, lights= lights,cameras=camera,materials=mater)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3wgeJgvsRoB1"
      },
      "outputs": [],
      "source": [
        "# Rasterization settings for silhouette rendering  \n",
        "sigma = 1e-5\n",
        "raster_settings_silhouette = RasterizationSettings(\n",
        "    image_size=256, \n",
        "    blur_radius=np.log(1. / 1e-4 - 1.)*sigma, \n",
        "    faces_per_pixel=20# Rasterization settings for silhouette rendering  \n",
        ")\n",
        "\n",
        "# Silhouette renderer \n",
        "renderer_silhouette = MeshRenderer(\n",
        "    rasterizer=MeshRasterizer(\n",
        "        cameras=camera, \n",
        "        raster_settings=raster_settings_silhouette\n",
        "    ),\n",
        "    shader=SoftSilhouetteShader()\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "o6BpJ2HtrRpL"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "  if True:#to debug cameras for trianing\n",
        "    with torch.no_grad():\n",
        "      images_,_ = renderer_train(mesh3, lights= lights,cameras=camera,materials=mater)\n",
        "      fig=plt.figure(figsize=(8,8))\n",
        "      plt.imshow(images_[0,:,:,:3].cpu())\n",
        "      plt.show()\n",
        "except Exception as e:\n",
        "  print (e)\n",
        "  pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SrsFhTdRm8si"
      },
      "source": [
        "### Sobel edges"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "92bjC3Pr4Ng9"
      },
      "outputs": [],
      "source": [
        "def sobel(x,abs=False,sq=True):\n",
        "    x=x.permute(0,3,1,2)\n",
        "    x=F.interpolate(x,size=(224,224))\n",
        "    assert(x.shape[1]<=4)\n",
        "    border1 = x[:, :, :-1] - x[:, :, 1:]\n",
        "    border1 = torch.cat([border1, x[:, :, :1] * 0], 2)  ##so square with extra 0 line\n",
        "    border2 = x[:, :, :, :-1] - x[:, :, :, 1:]\n",
        "    border2 = torch.cat([border2, x[:, :, :, :1] * 0], 3)\n",
        "    if abs:\n",
        "        border = torch.cat([border1.abs().mean(1).unsqueeze(1), border2.abs().mean(1).unsqueeze(1)], 1)##keep orientation of edge\n",
        "        return border\n",
        "    if sq:\n",
        "        border=(border1**2 + border2**2).mean(1).unsqueeze(1)#.sqrt()\n",
        "        with torch.no_grad():\n",
        "          m = border.mean()\n",
        "          s=border.std()\n",
        "        return (border-m)/s\n",
        "        #border = 1 - (-2 * (1e-5 + border).sqrt()).exp()  ##so no edge is 0, any edge quickly goes to 1\n",
        "        return border\n",
        "    border = border1+border2# torch.cat([border1, border2], 1)\n",
        "    return border\n",
        "\n",
        "def alignTS(a,b):\n",
        "      delta_bw_color = (sobel(a)-sobel(b)).abs().mean()\n",
        "      return delta_bw_color#faster"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s7w7dsnqCHzy"
      },
      "outputs": [],
      "source": [
        "def priorT(img):#img is B H W C\n",
        "        tv_h = ((img[:,:,1:] - img[:,:,:-1]).pow(2)).mean()\n",
        "        tv_w = ((img[:,1:,:] - img[:,:-1,:]).pow(2)).mean()    \n",
        "        return (tv_h + tv_w)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IUds2p6aRoB5"
      },
      "source": [
        "### directonal light and material - simplified"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nKJZqHspRoB6"
      },
      "outputs": [],
      "source": [
        "lparams = torch.cuda.FloatTensor(12).uniform_(-0.05,0.05)\n",
        "lparams=nn.Parameter(lparams*5)\n",
        "\n",
        "lparams2 = torch.cuda.FloatTensor(12).uniform_(-0.05,0.05)\n",
        "lparams2=nn.Parameter(lparams2*5)\n",
        "\n",
        "mparams = torch.cuda.FloatTensor(1,10).uniform_(-0.05,0.05)\n",
        "mparams=nn.Parameter(mparams*5)\n",
        "\n",
        "#all white -- just light enough to change colors\n",
        "def calcMaterials():  \n",
        "  sc=0.1\n",
        "  mater=Materials(shininess=torch.sigmoid(mparams[0,9])*999)\n",
        "  #mater=Materials(shininess=500)\n",
        "  mater=mater.cuda()#so shininess 450 default\n",
        "  return mater"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wTvrowjLRoB6"
      },
      "outputs": [],
      "source": [
        "def calcLights_(lparams,verb):\n",
        "    if True:\n",
        "        dif=  0.3+0.4*torch.sigmoid(lparams[None, 3:6])#range 0.3 to 0.7\n",
        "        amb=  0.15+0.3*torch.sigmoid(lparams[None, 6:9])#range 0.2 to.5\n",
        "        spec=  0.2+0.3*torch.sigmoid(lparams[None, 9:12])\n",
        "        \n",
        "        s=amb+dif+spec#-1e-1#TODO divide?\n",
        "        amb=amb/s\n",
        "        dif=dif/s\n",
        "        spec=spec/s\n",
        "    \n",
        "    if verb:\n",
        "            print (\"light\",amb,dif,spec)    \n",
        "    \n",
        "    lights=DirectionalLights(ambient_color=amb, diffuse_color=dif,\n",
        "            specular_color=spec, direction=((0, 5, 5), ), device=device)\n",
        "    return lights\n",
        "\n",
        "def calcLights(verb=False):\n",
        "  return calcLights_(lparams,verb)\n",
        "def calcLights2(verb=False):\n",
        "  return calcLights_(lparams2,verb)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## main deform optimize routines"
      ],
      "metadata": {
        "id": "6CORkeNd2pI1"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## calibration"
      ],
      "metadata": {
        "id": "_Udo3A1INjah"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#return mean and scale, size B13\n",
        "def calibrateF(wc,v2):\n",
        "    with torch.no_grad():\n",
        "        if len(wc.shape)==2:\n",
        "            wc=wc.unsqueeze(0)  \n",
        "            \n",
        "        if wc.shape[0] != v2.shape[0]:\n",
        "            print (wc.shape,v2.shape)\n",
        "            raise Exception(\"cam discrepancy\")\n",
        "\n",
        "        B=wc.shape[0]\n",
        "        M = torch.zeros(B,1,3).cuda()\n",
        "        S = torch.zeros(B,1,1).cuda()#isoscale\n",
        "        for b in range(B):#TODO make batch op\n",
        "            for c in range(3):\n",
        "                ma = v2[b,:,c].max()\n",
        "                mi = v2[b,:,c].min()\n",
        "                M[b,0,c]=(ma+mi)/2\n",
        "            #print (\"mami sanity\",wc[b,:,0].max()-wc[b,:,0].min(),wc[b,:,1].max()-wc[b,:,1].min())\n",
        "            S[b,0,0] = max(wc[b,:,0].max()-wc[b,:,0].min(),wc[b,:,1].max()-wc[b,:,1].min())\n",
        "    return (M,2/S)\n",
        "\n",
        "#max interval\n",
        "def MaMi0(x):\n",
        "      return x.max()-x.min()\n",
        "#mean\n",
        "def MaMi1(x):\n",
        "      return 0.5*(x.max()+x.min()).unsqueeze(0)\n",
        "\n",
        "def sampleOrient():\n",
        "  s=viewport_limit\n",
        "  buf= list(range(180-s,180+s)) + list(range(360-s,360+s))\n",
        "  return buf[np.random.randint(len(buf))]"
      ],
      "metadata": {
        "id": "-Zt6nVfRNm07"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### new mesh part transform"
      ],
      "metadata": {
        "id": "U0O6JeTr36FK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from pytorch3d.transforms import Transform3d\n",
        "#first version: 4 params; later add xyz rotate\n",
        "def trans(v,theta=None):\n",
        "  if len(v.shape) ==2:\n",
        "    print (\"weird unsqueeze\")\n",
        "    v=v.unsqueeze(0)#so BxNx3 tensor\n",
        "  #print (\"trans\",v.shape,theta.shape,\"device\",v.device,theta.device)\n",
        "  out=[]\n",
        "\n",
        "  if np.random.rand()<0.15:\n",
        "    print (\"theta\",theta[5:])\n",
        "\n",
        "  reg=0\n",
        "  for i in range(NClones):\n",
        "    x=v[:,i*NVOrig:(i+1)*NVOrig]\n",
        "    if i==0:\n",
        "      out.append(x)\n",
        "    else:\n",
        "      out.append(trans_(x,theta[i*5:i*5+5]))\n",
        "      reg += (out[0]-out[-1]).abs().mean()\n",
        "  return torch.cat(out,1),reg\n",
        "\n",
        "def trans_(v,theta):\n",
        "  #out=v*theta[0].exp() + theta[1:4].view(1,1,3)\n",
        "  if False:\n",
        "    t1 = Transform3d().scale(theta[0].exp()).translate(theta[1],theta[2],theta[3]).cuda()\n",
        "  elif True:\n",
        "    t1 = Transform3d().cuda()\n",
        "  else:\n",
        "    t1 = Transform3d().translate(theta[1],theta[2],theta[3]).cuda()\n",
        "    rot_y = RotateAxisAngle(theta[0]*120,'Z').cuda()\n",
        "    t1=t1.compose(rot_y)\n",
        "\n",
        "  rot_y = RotateAxisAngle(theta[4]*120,'Y').cuda()\n",
        "  t1=t1.compose(rot_y)\n",
        "  out = t1.transform_points(v)\n",
        "  return out"
      ],
      "metadata": {
        "id": "fD4dGEGd39mf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2JBEm6CSMwl0"
      },
      "source": [
        "### deform mesh"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rq8Tsc0BAaiH"
      },
      "outputs": [],
      "source": [
        "def render_mesh(poses, textures,scaleAugm=1,tScale=0,\\\n",
        "               deform_value=None,rcam=None,\\\n",
        "               cam_rescale=False,calibrate=None,mesh_out_only=False,global_orient=None,verb=False, randomCrop=False,orient=None):  \n",
        "  if rcam is None:\n",
        "    rcam=sampleRC()  \n",
        "  if rcam is None:\n",
        "    rcam,_=sampleRC()\n",
        "  num_views = len(rcam)\n",
        "\n",
        "  if False:\n",
        "    v2 = ORACLE.net(orig_v_template).unsqueeze(0) #deform    \n",
        "    v2=v2.expand(num_views,-1,-1)  \n",
        "    v2,reg_part=trans(v2,ORACLE.theta)\n",
        "  else:\n",
        "    v2=orig_v_template.unsqueeze(0) \n",
        "    v2,reg_part=trans(v2,ORACLE.theta)\n",
        "    v2 = ORACLE.net(v2.squeeze()).unsqueeze(0) \n",
        "    v2=v2.expand(num_views,-1,-1)  \n",
        "\n",
        "  ORACLE.reg_part=reg_part\n",
        "\n",
        "  if orient is None:#random yaxis rotate, azimuth\n",
        "    orient = []\n",
        "    for i in range(num_views):\n",
        "      orient+=[sampleOrient()]\n",
        "  rot_y = RotateAxisAngle(orient,'Y', device=device)\n",
        "  v2= rot_y.transform_points(v2)\n",
        "\n",
        "  faces=f2.expand(num_views,-1,-1)  \n",
        "\n",
        "  mesh3 = Meshes(v2,faces)\n",
        "  if calibrate is not None:\n",
        "        me=calibrate[0]\n",
        "        dx=calibrate[1]\n",
        "        v2_=(v2-me)\n",
        "        v2_=v2_*2/dx\n",
        "        mesh3 = Meshes(v2_,faces) \n",
        "  \n",
        "  if randomCrop or True:#subset of mesh for training rendered in window\n",
        "    #ix=np.random.randint(low=0, high=v2.shape[1], size=(60,))\n",
        "    ix = range(NVOrig)\n",
        "    camera_verts=v2[:,ix]\n",
        "  else:\n",
        "    camera_verts=v2\n",
        "\n",
        "  wc=rcam.transform_points(camera_verts)#mesh3.verts_packed() in ndc space already\n",
        "  if cam_rescale:\n",
        "    mea,sca=calibrateF(wc,camera_verts)\n",
        "    v2=(v2-mea)*sca\n",
        "    mesh3=Meshes(v2,faces)\n",
        "      \n",
        "  if mesh_out_only:\n",
        "    return mesh3\n",
        "\n",
        "  mesh3.textures=textures[0].extend(num_views)\n",
        "  mesh3.normal_map=textures[1].extend(num_views)\n",
        "  lights = calcLights()\n",
        "  images,_ = renderer_train(mesh3, lights=lights,cameras=rcam,materials=calcMaterials())#random camera sample --- but save in canonical cameras later\n",
        "  images=torch.clip(images,0,1)\n",
        "  return images,mesh3\n",
        "\n",
        "buf=[]#store losses"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eo6jSmsaLtXR"
      },
      "outputs": [],
      "source": [
        "bRanger=True\n",
        "if not bRanger:\n",
        "  optimizer=torch.optim.Adam\n",
        "else:\n",
        "  #!python -m pip install git+https://github.com/lessw2020/Ranger21.git\n",
        "  if True:\n",
        "    !git clone https://github.com/lessw2020/Ranger21.git\n",
        "    %cd Ranger21\n",
        "    !python -m pip install -e .\n",
        "    %cd ..\n",
        "  sys.path.append(\"/content/Ranger21/ranger21\")\n",
        "  from ranger21 import Ranger21 \n",
        "  optimizer=Ranger21##hmm, more complex initialize\n",
        "  #o=Ranger21()\n",
        "bUNET = False##complex shading of BW\n",
        "GA=1\n",
        "pfr=40#frequency to print\n",
        "lr=0.003*2#/GA\n",
        "from time import time"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9A-nPSkhRoB-"
      },
      "source": [
        "### text CLIP routines"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "z31y23176JiT"
      },
      "outputs": [],
      "source": [
        "text_other = '''incoherent, confusing, cropped, watermarks, grainy, pixellated, noisy'''#, degenerate limbs, weird human body\n",
        "\n",
        "\n",
        "import imageio\n",
        "import torchvision\n",
        "\n",
        "def getCLIPE_(s):\n",
        "  if s[-3:]=='jpg' or s[-3:]=='png':\n",
        "    cpath=\"%s/\"%(resourcePath)\n",
        "    plt.imshow(imageio.imread(cpath+s))\n",
        "    plt.show()\n",
        "    img_enc = (torch.nn.functional.interpolate(torch.tensor(imageio.imread(cpath+s)).unsqueeze(0).permute(0, 3, 1, 2), (224, 224)) / 255).cuda()[:,:3]\n",
        "    img_enc = clip_preprocess(img_enc)\n",
        "    img_enc = perceptor.encode_image(img_enc.cuda()).detach().clone()\n",
        "    print (\"img prompt\")\n",
        "    return img_enc\n",
        "  else:\n",
        "    #s+= \" rendered in Unreal Engine\"\n",
        "    tx = clip.tokenize(s)\n",
        "    t = perceptor.encode_text(tx.cuda()).detach().clone()\n",
        "    return t\n",
        "\n",
        "def getCLIPE(s):\n",
        "  t=[]\n",
        "  for to in s.split(','):\n",
        "    if len(to)==0:\n",
        "        continue\n",
        "    print (to)\n",
        "    t.append(getCLIPE_(to))\n",
        "  return t\n",
        "\n",
        "t2=[]#empty, hack\n",
        "print (\"no neg words\")\n",
        "    \n",
        "def getClipTargets(data):\n",
        "    clipTargets=[]\n",
        "    for t in data:\n",
        "        tc0=getCLIPE(t)\n",
        "        clipTargets.append(tc0)\n",
        "    return clipTargets\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LpFwcXk0JlzN"
      },
      "source": [
        "### multiscale code"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "55rsYo1ERoCC"
      },
      "outputs": [],
      "source": [
        "vert1 = v10\n",
        "print (vert1.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vEvnCcT_RoCD"
      },
      "outputs": [],
      "source": [
        "bUseMscale = HRL>0 \n",
        "def Scale_CLIP(m,prompts,color=0):\n",
        "    total=0\n",
        "    out=[]\n",
        "    rcam=sampleRC()    \n",
        "    v2=m.verts_padded()\n",
        "    \n",
        "    meshes=[]\n",
        "    for t in prompts:        \n",
        "        v2_=v2[:,ix_level0,:]\n",
        "        #print (\"mscale stuff\",v2_.shape,f10.shape,len(rcam))\n",
        "        mesh = Meshes(v2_,f10.expand(len(rcam),-1,-1))\n",
        "        mesh.textures=TexturesVertex(verts_features=v2_*0+color) #just text_vertex\n",
        "        imagesBW,_= renderer_train2(mesh, lights= calcLights(),cameras=rcam,\\\n",
        "                                   materials=calcMaterials(),sizeExtra=0)#.extend(num_views)\n",
        "\n",
        "        l1bw= lossClip(imagesBW.permute(0,3,1,2)[:,:3],t,[])\n",
        "        total +=l1bw\n",
        "        out.append(imagesBW[0].cpu())\n",
        "    return total,torch.cat(out,1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5k57IxmQK708"
      },
      "outputs": [],
      "source": [
        "def initParams():  \n",
        "    lsc=30\n",
        "    lparams.data.uniform_(-0.1*lsc,0.1*lsc)##does it work and change function?\n",
        "    lparams2.data.uniform_(-0.1*lsc,0.1*lsc)\n",
        "    mparams.data.uniform_(-0.1*lsc,0.1*lsc)\n",
        "\n",
        "    bw_shade=nn.Parameter(torch.cuda.FloatTensor(1,3).uniform_(-2,2))#sigmoid afterwards anyway\n",
        "    texNet = TextureL().cuda()\n",
        "    mlp = PlainDeform(orig_v_template.squeeze(),initGamma=initGamma,initPV=initPV).cuda()   #orig_v_template \n",
        "\n",
        "    del_v=list(mlp.parameters())\n",
        "    ORACLE.net = mlp\n",
        "    ORACLE.last=None\n",
        "    if True:\n",
        "        ORACLE.shaderI=UNet(3,c_in=4+2+2,c=channels_neural_shader).cuda()\n",
        "        recolor=ORACLE.shaderI#hack to use this instead\n",
        "    else:\n",
        "        recolor=nn.Linear(3,3)#dummy\n",
        "        ORACLE.shaderI=None\n",
        "    print (mlp)\n",
        "\n",
        "    theta=torch.zeros(NClones*5).uniform_(-0.3,0.3).cuda()\n",
        "    theta=nn.Parameter(theta*0.2)\n",
        "  \n",
        "    return bw_shade,texNet,mlp,del_v,recolor,theta "
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### get part rendered image"
      ],
      "metadata": {
        "id": "FIS0aWYceHm4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#center camera on mesh subpart and show it fully\n",
        "def getHead(mesh3,iHead,rcam=None):\n",
        "      if rcam is None:\n",
        "        rcam=sampleRC()\n",
        "\n",
        "      v2=mesh3.verts_padded()        \n",
        "\n",
        "      wc=rcam.transform_points(v2[:,iHead,:])\n",
        "      mea,sca=calibrateF(wc,v2[:,iHead,:])\n",
        "      v2=(v2-mea)*sca\n",
        "      #mesh3=Meshes(v2,faces)\n",
        "      \n",
        "      mesh4 = Meshes(v2,mesh3.faces_padded())\n",
        "      mesh4.textures=mesh3.textures\n",
        "      mesh4.normal_map=mesh3.normal_map\n",
        "\n",
        "      lights = calcLights()\n",
        "      imagesHead,_ = renderer_train(mesh4, lights=lights,cameras=rcam,materials=calcMaterials(),sizeExtra=0)\n",
        "      imagesHead=torch.clip(imagesHead,0,1)\n",
        "      return imagesHead,rcam"
      ],
      "metadata": {
        "id": "aH5-53oZbuFI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### main optimization step-routine"
      ],
      "metadata": {
        "id": "vaiLkIVUNKwe"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def vis_iter(tx,texn,images,images_full,imbw_r,crops,buf,text_0,text_0bw):\n",
        "    plt.figure(figsize=(20,10))\n",
        "    plt.subplot(1,2,1)\n",
        "    vis_nm = texn[0,:,:,:3].detach().cpu()\n",
        "    nr = torch.sqrt((vis_nm**2).sum(2))\n",
        "    vis_nm = vis_nm/nr.unsqueeze(2)#so norm 1, btw -1 and 1\n",
        "    vis_nm = vis_nm*0.5+0.5#so 0 to 1\n",
        "    plt.imshow(vis_nm)\n",
        "    plt.title(\"normal maps\")\n",
        "    plt.subplot(1,2,2)\n",
        "    plt.imshow(tx[0].detach().cpu().clip(0,1))\n",
        "    plt.title(\"textures\")\n",
        "    plt.show()\n",
        "\n",
        "    print (\"time\",time()-t0)\n",
        "    plt.figure(figsize=(4*8,8))\n",
        "    plt.subplot(1,4,1)\n",
        "    vis11=images[...,:3].detach().cpu()\n",
        "    vis11 = vis11.chunk(num_views)\n",
        "    vis11=torch.cat(vis11,1).squeeze(0)\n",
        "    plt.imshow(vis11)\n",
        "    plt.title(text_0)\n",
        "    plt.subplot(1,4,2)\n",
        "    try:\n",
        "      plt.imshow(images_full[0])\n",
        "    except Exception as e:\n",
        "      print (e)\n",
        "    plt.title(text_0)\n",
        "    plt.subplot(1,4,3)\n",
        "    vis11=imbw_r[...,:3].detach().cpu()\n",
        "    vis11 = vis11.chunk(num_views)\n",
        "    vis11=torch.cat(vis11,1).squeeze(0)\n",
        "    plt.imshow(vis11)\n",
        "    plt.title(text_0bw)\n",
        "    plt.subplot(1,4,4)\n",
        "    if len(crops)>0:\n",
        "      crops=torch.cat(crops,1)[0,:,:,:3]\n",
        "      plt.imshow(crops.detach().cpu())\n",
        "\n",
        "    plt.show()\n",
        "    \n",
        "    a=np.array(buf)\n",
        "    print (a.shape,buf[-1])\n",
        "    labels=[text_0,text_0bw,\"color parts\",\"black and white parts\"]\n",
        "    for i in range(4):\n",
        "      plt.plot(a[:,i],label=labels[i])\n",
        "    plt.title('CLIP')\n",
        "    plt.legend()\n",
        "    plt.show()\n",
        "\n",
        "    names = ['total','tex','sil']+['edge','lap','normal'] + ['deform','edgeEach',\"formTexture_align\",\"shader_regularize\"]\n",
        "    a=a[:,4:]\n",
        "    plt.figure(figsize=(24,3))\n",
        "    for i in range(a.shape[1]):\n",
        "      plt.subplot(1,a.shape[1],i+1)\n",
        "      plt.plot(a[:,i])\n",
        "      plt.title(names[i])\n",
        "    plt.show()"
      ],
      "metadata": {
        "id": "UHJPDM5TVUpQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wFkH1a6ulLH1",
        "scrolled": false
      },
      "outputs": [],
      "source": [
        "def runStep(z,t0=time(),OpPa=None,clipTargets=None,dataText=None):\n",
        "  bw_shade,texNet,mlp,del_v,recolor,opti,buf,imagesA = OpPa\n",
        "  tc0= clipTargets[0]\n",
        "  tc0b=clipTargets[1]\n",
        "  text_0 = dataText[0]\n",
        "  text_0bw = dataText[1]\n",
        "    \n",
        "  torch.cuda.empty_cache()\n",
        "  if z%GA==0:\n",
        "    opti.zero_grad()  #what is this? HACK MAy 10 inspect \n",
        "    ORACLE.shadereg=0*(tc0[0]).sum()\n",
        "\n",
        "  if True:\n",
        "    tx,texn=texNet()#.permute(0,2,3,1)\n",
        "    if z<3:\n",
        "      plt.imshow(tx[0].detach().cpu())\n",
        "      plt.show()\n",
        "    textures =TexturesUV(maps=tx,faces_uvs=[faces_uvs], verts_uvs=vt)\n",
        "    texturesn =TexturesUV(maps=texn,faces_uvs=[faces_uvs], verts_uvs=vt)\n",
        "\n",
        "  rcam=sampleRC()\n",
        "  images,mesh3=render_mesh(None,(textures,texturesn),\\\n",
        "                                     deform_value=0,rcam=rcam,cam_rescale=True,randomCrop=True)##hack: was True -- if False will also optimize pose, allowing to cheat but get rid of dynamics!!\n",
        "  #print (\"runstep done\")\n",
        "  deform_value = mlp.last\n",
        "  l1= lossClip(procClip(images),tc0,t2)\n",
        "  if z <20:\n",
        "      print (\"l1 loss\",l1.item(),len(tc0),tc0[0].shape,text_0)\n",
        "  reg_siam0=l1*0\n",
        "  \n",
        "  texturesBW=TexturesUV(maps=tx.detach()*0+torch.sigmoid(bw_shade),faces_uvs=[faces_uvs], verts_uvs=vt)\n",
        "  mesh3.textures=texturesBW.extend(num_views)\n",
        "  imagesBW,depthBW = renderer_train(mesh3, lights= calcLights2(),cameras=rcam,materials=calcMaterials(),sizeExtra=0)#random camera sample --- but save in canononical cameras later\n",
        "  mesh3.textures=textures\n",
        "  depthBW=imagesBW\n",
        "\n",
        "  l1bw= lossClip(procClip(imagesBW),tc0b,[])\n",
        "  imbw_r=imagesBW#.permute(0,3,1,2)[:,:3]\n",
        "  delta_bw_color = alignTS(imbw_r,images)\n",
        "    \n",
        "  if bUseMscale:# and False:\n",
        "      l1bw_scale,out_scale=Scale_CLIP(mesh3,[tc0b],torch.sigmoid(bw_shade))\n",
        "      l1bw=l1bw+l1bw_scale\n",
        "  \n",
        "\n",
        "  def getSubLoss(iHead,tc1,tc1b,sizeExtra=0):\n",
        "      mesh3.textures=textures.extend(num_views)\n",
        "      imagesHead,rcam=getHead(mesh3,iHead)\n",
        "      l1head=lossClip(procClip(imagesHead),tc1,t2)\n",
        "      mesh3.textures=texturesBW.extend(num_views)\n",
        "      imagesHead_bw,_=getHead(mesh3,iHead,rcam=rcam)\n",
        "      delta_bw_color3 = alignTS(imagesHead,imagesHead_bw)\n",
        "      silh_hbw=imagesHead_bw[:,:,:,3:4]                                        \n",
        "      l1head_bw = lossClip(procClip(imagesHead_bw),tc1b,[])\n",
        "      with torch.no_grad():\n",
        "          if imagesHead.shape[2]!=224:\n",
        "            imagesHead =F.interpolate(imagesHead.permute(0,3,1,2),size=(224,224)).permute(0,2,3,1)\n",
        "          crop1=torch.cat([imagesHead[:,:,:],imagesHead_bw[:,:,:]],2)      \n",
        "      return l1head,l1head_bw,crop1\n",
        "\n",
        "  loss_aux=l1*0\n",
        "  loss_auxbw=l1*0\n",
        "  crops=[]\n",
        "  for i in range(1,len(clipTargets)//2):\n",
        "    iHead=parts[i-1]\n",
        "    tc1=clipTargets[i*2]\n",
        "    tc1b=clipTargets[i*2+1]\n",
        "    lpart,lpartbw,crop1=getSubLoss(iHead,tc1,tc1b,sizeExtra=0)\n",
        "    loss_aux+=lpart\n",
        "    loss_auxbw+=lpartbw\n",
        "    crops.append(crop1)\n",
        "\n",
        "  mesh3.textures=textures\n",
        "  silh=images[:,:,:,3]##soft silhouette is differentiable.. \n",
        " \n",
        "  l_e = (mesh_edge_loss(mesh3))# -l_e_orig).abs()\n",
        "  l_l = l_e*0#(mesh_laplacian_smoothing(mesh3))#-l_l_orig).abs()\n",
        "  l_n = (mesh_normal_consistency(mesh3))#-l_n_orig).abs\n",
        "  l_orig = (deform_value**2).mean()\n",
        "  l3 =priorT(tx)\n",
        "\n",
        "  l5= ((silh.mean(2).mean(1)-fSilhTarget)**2).mean()\n",
        "   \n",
        "  gen_edge_len= mesh_edge_len(mesh3)[:edge_len.shape[0]]  \n",
        "  loss_edge_len = ((wel*(edge_len-gen_edge_len ))**2).mean()\n",
        "  loss = l1bw+ l1+1e-3*l3+fSilhStrength *l5+\\\n",
        "  4e0*(1e-1*l_e+50*0.1*l_n+0*l_l)+fPenaltyDeformation*l_orig+fPenaltyDeformation_e*loss_edge_len \n",
        "\n",
        "  loss+=loss_aux+loss_auxbw\n",
        "  loss += 1e-1*(delta_bw_color)\n",
        "  loss += fShader_reg*ORACLE.shadereg\n",
        "  \n",
        "  loss +=fPartRegularize*ORACLE.reg_part\n",
        "  print (\"reg part\",ORACLE.reg_part)\n",
        "\n",
        "  loss.backward()\n",
        "  buf.append([l1.item(),l1bw.item(),loss_aux.item(),loss_auxbw.item(),loss.item(),l3.item(),silh.mean().item(),l_e.item(),l_l.item(),l_n.item(),l_orig.item(),loss_edge_len.item(),\\\n",
        "              delta_bw_color.item(),ORACLE.shadereg.item()])\n",
        "  \n",
        "  if z<=5 or z%10==0:\n",
        "    print (z,\"total\",loss.item(),\"texsm\",l3.item(),\"silh\",silh.mean().item(),\\\n",
        "       \"mesh\",l_e.item(),l_l.item(),l_n.item(),\"off-mesh\",l_orig.item())\n",
        "    \n",
        "    print (\"dbwcoolor \",delta_bw_color.item())\n",
        "    print (\"clip losses\",l1.item(),l1bw.item())\n",
        "    print (\"shader reg\",ORACLE.shadereg.item())\n",
        "    print (\"\")\n",
        "    \n",
        "  best =loss.item()\n",
        "  if z%GA ==GA-1:\n",
        "    opti.step()\n",
        "    if z%1==0:#save training frame for output animation\n",
        "      with torch.no_grad():       \n",
        "          target_cameras =getCamP(R=R2[:1], T=T2[:1])#choose 1 camera, but larger res\n",
        "          mesh3=render_mesh(None,(textures,texturesn),cam_rescale=True,mesh_out_only=True,orient=0,rcam=target_cameras)\n",
        "          v =mesh3.verts_packed()\n",
        "          mesh3.textures=textures#.extend(1)\n",
        "          mesh3.normal_map=texturesn#.extend(1)\n",
        "          ci = (z//GA)%R2.shape[0]#iterate cameras, rotate during saving\n",
        "          target_cameras =getCamP(R=R2[ci:ci+1], T=T2[ci:ci+1])#choose 1 camera, but larger res\n",
        "          lights = calcLights()\n",
        "          images_full = renderer(mesh3, lights= lights,cameras=target_cameras,materials=calcMaterials()).clamp(0,1).cpu()[:,:,:,:3]\n",
        "          imagesA.append(images_full.half())#\n",
        "    \n",
        "  if z%pfr==0 or z <5:\n",
        "    vis_iter(tx,texn,images,images_full,imbw_r,crops,buf,text_0,text_0bw)\n",
        "  return images,imagesBW"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OBGxvhXA5jnI"
      },
      "source": [
        "### final video animation routines"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ozn4RQm78a9N",
        "scrolled": false
      },
      "outputs": [],
      "source": [
        "def getCalibrate(textures,texturesn):\n",
        "    with torch.no_grad():##just any camera for roughly centering mesh\n",
        "      mesh3=render_mesh(None,( textures,texturesn),mesh_out_only=True)\n",
        "      ci=0\n",
        "      rcam =getCamP(R=R2[ci:ci+1], T=T2[ci:ci+1])\n",
        "      camera_verts=mesh3.verts_packed()[:NVOrig]#v2\n",
        "      wc=rcam.transform_points(camera_verts).unsqueeze(0)#if 1 camera otherwise jusdt 2d tensor\n",
        "      v2=mesh3.verts_padded()[:,:NVOrig]\n",
        "      print (v2.shape,v2.max(),v2.min(),\"wc\",wc.shape,wc.max(),wc.min())\n",
        "      dx= MaMi0(wc[:,:,0])\n",
        "      dy= MaMi0(wc[:,:,1])\n",
        "      me0=MaMi1(v2[:,:,0])\n",
        "      me1=MaMi1(v2[:,:,1])\n",
        "      me2=MaMi1(v2[:,:,2])\n",
        "      ##TODO average ovrr multiple frames\n",
        "      me=torch.cat([me0,me1+0,me2]).view(1,1,3)      \n",
        "      calibrate=(me,max(dx,dy)+0.15)\n",
        "      return calibrate\n",
        "    \n",
        "def save_frames(textures,texturesn,calibrate =None,Ta=None,Sa=None,Ae=None,rev=False):\n",
        "    !rm frames2/*\n",
        "    if calibrate is None:\n",
        "        calibrate= getCalibrate(textures,texturesn)\n",
        "        \n",
        "    corig = calibrate\n",
        "        \n",
        "    imagesA=[]\n",
        "    print (\"starting buf\",len(imagesA))\n",
        "    ci=0\n",
        "    T=360#time length of output video\n",
        "    if True:  \n",
        "      for t in range(T):#\n",
        "          if Ta is not None:\n",
        "              m=corig[0]*1\n",
        "              s = corig[1]*1  \n",
        "              alfa=t/float(T-1)#from 0 to 1\n",
        "              if rev:\n",
        "                    alfa=1-alfa\n",
        "              alfa = alfa**Ae#so more aggressive\n",
        "              m[...,1]+= alfa*Ta #Y axis towards gead\n",
        "              calibrate =(m,s/(1+Sa*alfa))\n",
        "          with torch.no_grad():\n",
        "            ci = (ci+1)%R2.shape[0]#iterate cameras\n",
        "            target_cameras=getCamP(R=R2[ci:ci+1], T=T2[ci:ci+1])\n",
        "            \n",
        "            mesh3=render_mesh(None,(textures,texturesn),calibrate=calibrate,mesh_out_only=True,orient=t-60,rcam=target_cameras)\n",
        "            mesh3.normal_map=texturesn#last call from training\n",
        "            \n",
        "            if True:\n",
        "              mesh3.textures=textures                 \n",
        "              images_ = renderer(mesh3, lights=calcLights(),materials=calcMaterials(),cameras=target_cameras)\n",
        "              if False:\n",
        "                 images_=bwWindow(mesh3,images_,target_cameras,calcLights2(),z=ztarget)     \n",
        "            else:\n",
        "              mesh3.textures=texturesBW\n",
        "              imagesBW = renderer(mesh3, lights=calcLights2(),cameras=target_cameras,materials=calcMaterials())\n",
        "              images_=imagesBW\n",
        "\n",
        "            imagesA.append(images_.clamp(0,1).cpu().half())\n",
        "\n",
        "            if t%60==0:\n",
        "                plt.imshow(images_[0,:,:,:3].detach().cpu())\n",
        "                plt.title(\"motion step %s\"%(t))\n",
        "                plt.show()\n",
        "    return imagesA,mesh3"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9XuRr0SZRoCJ"
      },
      "outputs": [],
      "source": [
        "def director(final_name,textures,texturesn,Ta = 0.5,Sa=1.35,Ae=0.4,rev=True,imagesA=None):\n",
        "    if imagesA is not None:\n",
        "        !rm frames2/*\n",
        "        mesh3 = None\n",
        "    else:\n",
        "        imagesA,mesh3=save_frames(textures,texturesn,Ta = Ta,Sa=Sa,Ae=Ae,rev=rev)\n",
        "        for i in range(len(imagesA)):\n",
        "          im=imagesA[i]\n",
        "          if im.shape[3]==4:\n",
        "            imagesA[i] = imagesA[i][:,:,:,:3]\n",
        "    saveImageSet(mesh3,imagesA,\"t\")\n",
        "    \n",
        "    mname = \"movier4flat.mov\"\n",
        "    delCmd = \"rm %s\"%(mname)\n",
        "    os.system(delCmd)\n",
        "    \n",
        "    mCmd = 'ffmpeg -framerate 24 -r 24 -i \"%s'%(\"frames2\") +  '/%d.jpg\"'\n",
        "    mCmd += \" -c:v libx264 -crf 24 -pix_fmt yuv420p %s\"%(mname)\n",
        "    os.system(mCmd)    \n",
        "    shutil.copy(mname,final_name)\n",
        "    return mesh3"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1q3CCBeRRoCJ"
      },
      "source": [
        "#  optimization loop"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Vzb8SzTlRoCK",
        "scrolled": false
      },
      "outputs": [],
      "source": [
        "#@title Create!\n",
        "from IPython.display import clear_output\n",
        "for nText in range(2):\n",
        "    clear_output(wait=True)\n",
        "    dataText,descriptor=getTexts()##\n",
        "    \n",
        "    clipTargets=getClipTargets(dataText)\n",
        "    bw_shade,texNet,mlp,del_v,recolor,theta = initParams()\n",
        "    ORACLE.theta=theta\n",
        "    lr=0.003\n",
        "    opti = optimizer([{'params': [lparams,lparams2,mparams,bw_shade], 'lr':lr},\\\n",
        "                    {'params': del_v+ list(recolor.parameters()), 'lr':4e-3},\\\n",
        "                    {'params': theta, 'lr':1e-1},\\\n",
        "                    {'params': list(texNet.parameters()), 'lr':4e-1}], lr=lr,num_epochs =1,num_batches_per_epoch=GA*optisteps)\n",
        "    OpPa = (bw_shade,texNet,mlp,del_v,recolor,opti,[],[])  \n",
        "    \n",
        "    t0 = time()\n",
        "    for z in range(optisteps*GA):\n",
        "        images,imagesBW =runStep(z,t0,OpPa,clipTargets,dataText)#output just for vis\n",
        "        \n",
        "    ##when moved to function -- some stuff needs to be prepared again\n",
        "    ci=0\n",
        "    target_cameras =getCamP(R=R2[ci:ci+1], T=T2[ci:ci+1])\n",
        "    tx,texn=texNet()#.permute(0,2,3,1)\n",
        "    textures =TexturesUV(maps=tx,faces_uvs=[faces_uvs], verts_uvs=vt)\n",
        "    texturesn =TexturesUV(maps=texn,faces_uvs=[faces_uvs], verts_uvs=vt)\n",
        "    texturesBW=TexturesUV(maps=tx.detach()*0+torch.sigmoid(bw_shade),faces_uvs=[faces_uvs], verts_uvs=vt)\n",
        "    ORACLE.last=mlp.last\n",
        "    final_name=\"%s/CanonicS_%s_sym%s_NE%d.mov\"%(vid_dir,descriptor.replace('.','_'),mlp.bSYM,NE,)\n",
        "\n",
        "    if(create_video):\n",
        "      #video 1: rotate in full size visible\n",
        "      mesh3=director(final_name,textures,texturesn,Ta = None,Sa=None,Ae=0.4,rev=True)\n",
        "      #video 2: rotate and zoom from top\n",
        "      #mesh3=director(final_name.replace(\"CanonicS\",\"movZoomS\"),textures,texturesn,Ta = 0.55,Sa=1.35,Ae=0.75,rev=True)\n",
        "      #video 3: show evolution of trained mesh\n",
        "      _=director(final_name.replace(\"CanonicS\",\"iteration\"),textures,texturesn,imagesA=OpPa[-1])\n",
        "\n",
        "\n",
        "    if(bsave_obj):\n",
        "      save_obj('%s.obj'%(descriptor,),verts=mesh3.verts_packed(),faces=mesh3.faces_packed(),\\\n",
        "                 verts_uvs = vt.squeeze(),faces_uvs=faces_uvs.squeeze(),texture_map=textures.maps_padded().squeeze())\n",
        "      shutil.move(descriptor+'.obj', f'{meshDirPath}/{descriptor}.obj')\n",
        "\n",
        "    if(save_texture):\n",
        "      img=np.uint8(tx[0].cpu().detach().clip(0,1)*255)#tx and not tex due to sigmoid and other mods\n",
        "      imageio.imwrite(descriptor+'texture.png',img)\n",
        "      shutil.move(descriptor+'texture.png', f'{meshDirPath}/{descriptor}_texture.png')"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "texNet.tex.shape"
      ],
      "metadata": {
        "id": "mwd4ymkyEzSE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Interactive 3d Model\n"
      ],
      "metadata": {
        "id": "LwhxdpP5hgFY"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oPzGhe61mLI9"
      },
      "outputs": [],
      "source": [
        "if(show_3d_interactive_runall):\n",
        "  #mesh3.textures=texturesBW\n",
        "  fig = plot_scene({\n",
        "      \"text_to_add\": {\n",
        "          \"mean optimized\": mesh3\n",
        "      }\n",
        "  })\n",
        "  fig.show()"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "KeWlyoXCwzb5",
        "LYPgKJhfre8D"
      ],
      "name": "open_ClipMatrix_ver0.8.ipynb",
      "provenance": [],
      "private_outputs": true,
      "machine_shape": "hm",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "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.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}