[
  {
    "path": "LICENSE",
    "content": "                                 Apache License\n                           Version 2.0, January 2004\n                        http://www.apache.org/licenses/\n\n   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n   1. Definitions.\n\n      \"License\" shall mean the terms and conditions for use, reproduction,\n      and distribution as defined by Sections 1 through 9 of this document.\n\n      \"Licensor\" shall mean the copyright owner or entity authorized by\n      the copyright owner that is granting the License.\n\n      \"Legal Entity\" shall mean the union of the acting entity and all\n      other entities that control, are controlled by, or are under common\n      control with that entity. For the purposes of this definition,\n      \"control\" means (i) the power, direct or indirect, to cause the\n      direction or management of such entity, whether by contract or\n      otherwise, or (ii) ownership of fifty percent (50%) or more of the\n      outstanding shares, or (iii) beneficial ownership of such entity.\n\n      \"You\" (or \"Your\") shall mean an individual or Legal Entity\n      exercising permissions granted by this License.\n\n      \"Source\" form shall mean the preferred form for making modifications,\n      including but not limited to software source code, documentation\n      source, and configuration files.\n\n      \"Object\" form shall mean any form resulting from mechanical\n      transformation or translation of a Source form, including but\n      not limited to compiled object code, generated documentation,\n      and conversions to other media types.\n\n      \"Work\" shall mean the work of authorship, whether in Source or\n      Object form, made available under the License, as indicated by a\n      copyright notice that is included in or attached to the work\n      (an example is provided in the Appendix below).\n\n      \"Derivative Works\" shall mean any work, whether in Source or Object\n      form, that is based on (or derived from) the Work and for which the\n      editorial revisions, annotations, elaborations, or other modifications\n      represent, as a whole, an original work of authorship. For the purposes\n      of this License, Derivative Works shall not include works that remain\n      separable from, or merely link (or bind by name) to the interfaces of,\n      the Work and Derivative Works thereof.\n\n      \"Contribution\" shall mean any work of authorship, including\n      the original version of the Work and any modifications or additions\n      to that Work or Derivative Works thereof, that is intentionally\n      submitted to Licensor for inclusion in the Work by the copyright owner\n      or by an individual or Legal Entity authorized to submit on behalf of\n      the copyright owner. For the purposes of this definition, \"submitted\"\n      means any form of electronic, verbal, or written communication sent\n      to the Licensor or its representatives, including but not limited to\n      communication on electronic mailing lists, source code control systems,\n      and issue tracking systems that are managed by, or on behalf of, the\n      Licensor for the purpose of discussing and improving the Work, but\n      excluding communication that is conspicuously marked or otherwise\n      designated in writing by the copyright owner as \"Not a Contribution.\"\n\n      \"Contributor\" shall mean Licensor and any individual or Legal Entity\n      on behalf of whom a Contribution has been received by Licensor and\n      subsequently incorporated within the Work.\n\n   2. Grant of Copyright License. Subject to the terms and conditions of\n      this License, each Contributor hereby grants to You a perpetual,\n      worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n      copyright license to reproduce, prepare Derivative Works of,\n      publicly display, publicly perform, sublicense, and distribute the\n      Work and such Derivative Works in Source or Object form.\n\n   3. Grant of Patent License. Subject to the terms and conditions of\n      this License, each Contributor hereby grants to You a perpetual,\n      worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n      (except as stated in this section) patent license to make, have made,\n      use, offer to sell, sell, import, and otherwise transfer the Work,\n      where such license applies only to those patent claims licensable\n      by such Contributor that are necessarily infringed by their\n      Contribution(s) alone or by combination of their Contribution(s)\n      with the Work to which such Contribution(s) was submitted. If You\n      institute patent litigation against any entity (including a\n      cross-claim or counterclaim in a lawsuit) alleging that the Work\n      or a Contribution incorporated within the Work constitutes direct\n      or contributory patent infringement, then any patent licenses\n      granted to You under this License for that Work shall terminate\n      as of the date such litigation is filed.\n\n   4. Redistribution. You may reproduce and distribute copies of the\n      Work or Derivative Works thereof in any medium, with or without\n      modifications, and in Source or Object form, provided that You\n      meet the following conditions:\n\n      (a) You must give any other recipients of the Work or\n          Derivative Works a copy of this License; and\n\n      (b) You must cause any modified files to carry prominent notices\n          stating that You changed the files; and\n\n      (c) You must retain, in the Source form of any Derivative Works\n          that You distribute, all copyright, patent, trademark, and\n          attribution notices from the Source form of the Work,\n          excluding those notices that do not pertain to any part of\n          the Derivative Works; and\n\n      (d) If the Work includes a \"NOTICE\" text file as part of its\n          distribution, then any Derivative Works that You distribute must\n          include a readable copy of the attribution notices contained\n          within such NOTICE file, excluding those notices that do not\n          pertain to any part of the Derivative Works, in at least one\n          of the following places: within a NOTICE text file distributed\n          as part of the Derivative Works; within the Source form or\n          documentation, if provided along with the Derivative Works; or,\n          within a display generated by the Derivative Works, if and\n          wherever such third-party notices normally appear. The contents\n          of the NOTICE file are for informational purposes only and\n          do not modify the License. You may add Your own attribution\n          notices within Derivative Works that You distribute, alongside\n          or as an addendum to the NOTICE text from the Work, provided\n          that such additional attribution notices cannot be construed\n          as modifying the License.\n\n      You may add Your own copyright statement to Your modifications and\n      may provide additional or different license terms and conditions\n      for use, reproduction, or distribution of Your modifications, or\n      for any such Derivative Works as a whole, provided Your use,\n      reproduction, and distribution of the Work otherwise complies with\n      the conditions stated in this License.\n\n   5. Submission of Contributions. Unless You explicitly state otherwise,\n      any Contribution intentionally submitted for inclusion in the Work\n      by You to the Licensor shall be under the terms and conditions of\n      this License, without any additional terms or conditions.\n      Notwithstanding the above, nothing herein shall supersede or modify\n      the terms of any separate license agreement you may have executed\n      with Licensor regarding such Contributions.\n\n   6. Trademarks. This License does not grant permission to use the trade\n      names, trademarks, service marks, or product names of the Licensor,\n      except as required for reasonable and customary use in describing the\n      origin of the Work and reproducing the content of the NOTICE file.\n\n   7. Disclaimer of Warranty. Unless required by applicable law or\n      agreed to in writing, Licensor provides the Work (and each\n      Contributor provides its Contributions) on an \"AS IS\" BASIS,\n      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n      implied, including, without limitation, any warranties or conditions\n      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n      PARTICULAR PURPOSE. You are solely responsible for determining the\n      appropriateness of using or redistributing the Work and assume any\n      risks associated with Your exercise of permissions under this License.\n\n   8. Limitation of Liability. In no event and under no legal theory,\n      whether in tort (including negligence), contract, or otherwise,\n      unless required by applicable law (such as deliberate and grossly\n      negligent acts) or agreed to in writing, shall any Contributor be\n      liable to You for damages, including any direct, indirect, special,\n      incidental, or consequential damages of any character arising as a\n      result of this License or out of the use or inability to use the\n      Work (including but not limited to damages for loss of goodwill,\n      work stoppage, computer failure or malfunction, or any and all\n      other commercial damages or losses), even if such Contributor\n      has been advised of the possibility of such damages.\n\n   9. Accepting Warranty or Additional Liability. While redistributing\n      the Work or Derivative Works thereof, You may choose to offer,\n      and charge a fee for, acceptance of support, warranty, indemnity,\n      or other liability obligations and/or rights consistent with this\n      License. However, in accepting such obligations, You may act only\n      on Your own behalf and on Your sole responsibility, not on behalf\n      of any other Contributor, and only if You agree to indemnify,\n      defend, and hold each Contributor harmless for any liability\n      incurred by, or claims asserted against, such Contributor by reason\n      of your accepting any such warranty or additional liability.\n\n   END OF TERMS AND CONDITIONS\n\n   APPENDIX: How to apply the Apache License to your work.\n\n      To apply the Apache License to your work, attach the following\n      boilerplate notice, with the fields enclosed by brackets \"[]\"\n      replaced with your own identifying information. (Don't include\n      the brackets!)  The text should be enclosed in the appropriate\n      comment syntax for the file format. We also recommend that a\n      file or class name and description of purpose be included on the\n      same \"printed page\" as the copyright notice for easier\n      identification within third-party archives.\n\n   Copyright [yyyy] [name of copyright owner]\n\n   Licensed under the Apache License, Version 2.0 (the \"License\");\n   you may not use this file except in compliance with the License.\n   You may obtain a copy of the License at\n\n       http://www.apache.org/licenses/LICENSE-2.0\n\n   Unless required by applicable law or agreed to in writing, software\n   distributed under the License is distributed on an \"AS IS\" BASIS,\n   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n   See the License for the specific language governing permissions and\n   limitations under the License.\n"
  },
  {
    "path": "README.md",
    "content": "# InverseRenderNet: Learning single image inverse rendering\n\n***!! Check out our new work InverseRenderNet++ [paper](https://arxiv.org/abs/2102.06591) and [code](https://github.com/YeeU/InverseRenderNet_v2), which improves the inverse rendering results and shadow handling.***\n\nThis is the implementation of the paper \"InverseRenderNet: Learning single image inverse rendering\". The model is implemented in tensorflow.\n\nIf you use our code, please cite the following paper:\n\n    @inproceedings{yu19inverserendernet,\n        title={InverseRenderNet: Learning single image inverse rendering},\n        author={Yu, Ye and Smith, William AP},\n        booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},\n        year={2019}\n    }\n\n## Evaluation\n\n#### Dependencies\nTo run our evaluation code, please create your environment based on following dependencies:\n\n    tensorflow 1.12.0\n    python 3.6\n    skimage\n    cv2\n    numpy\n\n#### Pretrained model\n* Download our pretrained model from: [Link](https://drive.google.com/uc?export=download&id=1VKeByvprmWWXSig-7-fxfXs3KA-HG_-P)\n* Unzip the downloaded file \n* Make sure the model files are placed in a folder named \"irn_model\"\n\n\n#### Test on demo image\nYou can perform inverse rendering on random RGB image by our pretrained model. To run the demo code, you need to specify the path to pretrained model, path to RGB image and corresponding mask which masked out sky in the image. The mask can be generated by PSPNet, which you can find on https://github.com/hszhao/PSPNet. Finally inverse rendering results will be saved to the output folder named by your argument.\n\n```bash\npython3 test_demo.py --model /PATH/TO/irn_model --image demo.jpg --mask demo_mask.jpg --output test_results\n```\n\n\n#### Test on IIW\n* IIW dataset should be downloaded firstly from http://opensurfaces.cs.cornell.edu/publications/intrinsic/#download \n\n* Run testing code where you need to specify the path to model and IIW data:\n```bash\npython3 test_iiw.py --model /PATH/TO/irn_model --iiw /PATH/TO/iiw-dataset\n```\n\n## Training\n\n#### Train from scratch\nThe training for InverseRenderNet contains two stages: pre-train and self-train.\n* To begin with pre-train stage, you need to use training command specifying option `-m` to `pre-train`. \n* After finishing pre-train stage, you can run self-train by specifying option `-m` to `self-train`. \n\nIn addition, you can control the size of batch in training, and the path to training data should be specified.\n\nAn example for training command:\n```bash\npython3 train.py -n 2 -p Data -m pre-train\n```\n\n#### Data for training\nTo directly use our code for training, you need to pre-process the training data to match the data format as shown in examples in `Data` folder. \n\nIn particular, we pre-process the data before training, such that five images with great overlaps are bundled up into one mini-batch, and images are resized and cropped to a shape of 200 * 200 pixels. Along with input images associated depth maps, camera parameters, sky masks and normal maps are stored in the same mini-batch. For efficiency, every mini-batch containing all training elements for 5 involved images are saved as a pickle file. While training the data feeding thread directly load each mini-batch from corresponding pickle file.\n\n\n\n\n\n"
  },
  {
    "path": "model/SfMNet.py",
    "content": "import importlib\nimport tensorflow as tf\nimport numpy as np\nimport tensorflow.contrib.layers as layers\n\ndef SfMNet(inputs, height, width, name='', n_layers=12, n_pools=2, is_training=True, depth_base=64):\n\tconv_layers = np.int32(n_layers/2) -1\n\tdeconv_layers = np.int32(n_layers/2)\n\t# number of layers before perform pooling\n\tnlayers_befPool = np.int32(np.ceil((conv_layers-1)/n_pools)-1)\n\n\tmax_depth = 512\n\n\tif depth_base*2**n_pools < max_depth:\n\t\ttail = conv_layers - nlayers_befPool*n_pools\n\n\n\t\ttail_deconv = deconv_layers - nlayers_befPool*n_pools\n\telse:\n\t\tmaxNum_pool = np.log2(max_depth / depth_base)\n\t\ttail = np.int32(conv_layers - nlayers_befPool * maxNum_pool)\n\t\ttail_deconv = np.int32(deconv_layers - nlayers_befPool * maxNum_pool)\n\n\tf_in_conv = [3] + [np.int32(depth_base*2**(np.ceil(i/nlayers_befPool)-1)) for i in range(1, conv_layers-tail+1)] + [np.int32(depth_base*2**maxNum_pool) for i in range(conv_layers-tail+1, conv_layers+1)]\n\tf_out_conv = [64] + [np.int32(depth_base*2**(np.floor(i/nlayers_befPool))) for i in range(1, conv_layers-tail+1)] + [np.int32(depth_base*2**maxNum_pool) for i in range(conv_layers-tail+1, conv_layers+1)]\n\n\tf_in_deconv = f_out_conv[:0:-1] + [64]\n\tf_out_amDeconv = f_in_conv[:0:-1] + [3]\n\tf_out_MaskDeconv = f_in_conv[:0:-1] + [2]\n\tf_out_nmDeconv = f_in_conv[:0:-1] + [2]\n\n\n\n\tbatch_norm_params = {'decay':0.9, 'center':True, 'scale':True, 'epsilon':1e-4, 'param_initializers':{'beta_initializer':tf.zeros_initializer(),'gamma_initializer':tf.ones_initializer(),'moving_variance_initializer':tf.ones_initializer(),'moving_average_initializer':tf.zeros_initializer()}, 'param_regularizers':{'beta_regularizer':None,'gamma_regularizer':layers.l2_regularizer(scale=1e-5)},'is_training':is_training,'trainable':is_training} \n\n\t### contractive conv_layer block\n\tconv_out = inputs\n\tconv_out_list = []\n\tfor i,f_in,f_out in zip(range(1,conv_layers+2),f_in_conv,f_out_conv):\n\t\tscope = name+'conv'+str(i)\n\n\t\tif np.mod(i-1,nlayers_befPool)==0 and i<=n_pools*nlayers_befPool+1 and i != 1:\n\t\t\tconv_out_list.append(conv_out)\n\t\t\tconv_out = layers.conv2d(conv_out,num_outputs=f_out,kernel_size=[3,3],stride=[1,1],padding='SAME',normalizer_fn=layers.batch_norm, normalizer_params=batch_norm_params,weights_initializer=tf.random_normal_initializer(mean=0,stddev=np.sqrt(2/9/f_in)),weights_regularizer=layers.l2_regularizer(scale=1e-5),biases_initializer=None,scope=scope, trainable=is_training)\n\t\t\tconv_out = tf.nn.max_pool(conv_out, ksize=[1,2,2,1], strides=[1,2,2,1], padding='SAME')\n\n\t\telse:\n\n\t\t\tconv_out = layers.conv2d(conv_out,num_outputs=f_out,kernel_size=[3,3],stride=[1,1],padding='SAME',normalizer_fn=layers.batch_norm, normalizer_params=batch_norm_params,weights_initializer=tf.random_normal_initializer(mean=0,stddev=np.sqrt(2/9/f_in)),weights_regularizer=layers.l2_regularizer(scale=1e-5),biases_initializer=None,scope=scope, trainable=is_training)\n\n\n\t### expanding deconv_layer block succeeding conv_layer block\n\tam_deconv_out = conv_out\n\tfor i,f_in,f_out in zip(range(1,deconv_layers+1),f_in_deconv,f_out_amDeconv):\n\t\tscope = name+'am/am_deconv'+str(i)\n\n\t\t# expand resolution every after nlayers_befPool deconv_layer\n\t\tif np.mod(i,nlayers_befPool)==0 and i<=n_pools*nlayers_befPool:\n\t\t\twith tf.variable_scope(scope):\n\t\t\t\tW = tf.get_variable(regularizer=layers.l2_regularizer(scale=1e-5),initializer=get_bilinear_filter([3,3,f_out,f_in],2),shape=[3,3,f_out,f_in],name='filter', trainable=is_training)\t\n\t\t\t\t# import ipdb; ipdb.set_trace()\n\t\t\t\t# attach previous convolutional output to upsampling/deconvolutional output\n\t\t\t\ttmp = conv_out_list[-np.int32(i/nlayers_befPool)]\n\t\t\t\toutput_shape = tf.shape(tmp)\n\t\t\t\tam_deconv_out = tf.nn.conv2d_transpose(am_deconv_out,filter=W,output_shape=output_shape,strides=[1,2,2,1],padding='SAME')\n\t\t\t\tam_deconv_out = layers.batch_norm(scope=scope,activation_fn=tf.nn.relu,inputs=am_deconv_out,decay=0.9, center=True, scale=True, param_initializers={'beta_initializer':tf.zeros_initializer(),'gamma_initializer':tf.ones_initializer(),'moving_variance_initializer':tf.ones_initializer(),'moving_average_initializer':tf.zeros_initializer()}, param_regularizers={'beta_regularizer':None,'gamma_regularizer':layers.l2_regularizer(scale=1e-5)},is_training=is_training,trainable=is_training)\n\n\n\t\t\ttmp = layers.conv2d(tmp,num_outputs=f_out,kernel_size=[3,3],stride=[1,1],padding='SAME',normalizer_fn=layers.batch_norm, normalizer_params=batch_norm_params,weights_initializer=tf.random_normal_initializer(mean=0,stddev=np.sqrt(2/9/f_in)),weights_regularizer=layers.l2_regularizer(scale=1e-5),biases_initializer=None,scope=scope,trainable=is_training)\n\t\t\tam_deconv_out = tmp + am_deconv_out\n\n\n\t\telif i==deconv_layers:\n\t\t\tam_deconv_out = layers.conv2d(am_deconv_out,num_outputs=f_out,kernel_size=[3,3],stride=[1,1],padding='SAME',normalizer_fn=None,activation_fn=None,weights_initializer=tf.random_normal_initializer(mean=0,stddev=np.sqrt(2/9/f_in)),weights_regularizer=layers.l2_regularizer(scale=1e-5),scope=scope,trainable=is_training)\n\n\n\t\telse:\n\t\t\tam_deconv_out = layers.conv2d(am_deconv_out,num_outputs=f_out,kernel_size=[3,3],stride=[1,1],padding='SAME',normalizer_fn=layers.batch_norm, normalizer_params=batch_norm_params,weights_initializer=tf.random_normal_initializer(mean=0,stddev=np.sqrt(2/9/f_in)),weights_regularizer=layers.l2_regularizer(scale=1e-5),biases_initializer=None,scope=scope,trainable=is_training)\n\n\n\n\t### deconvolution net for nm estimates\n\tnm_deconv_out = conv_out\n\tfor i,f_in,f_out in zip(range(1,deconv_layers+1),f_in_deconv,f_out_nmDeconv):\n\t\tscope = name+'nm/nm'+str(i)\n\n\t\t# expand resolution every after nlayers_befPool deconv_layer\n\t\tif np.mod(i,nlayers_befPool)==0 and i<=n_pools*nlayers_befPool:\n\t\t\twith tf.variable_scope(scope):\n\t\t\t\tW = tf.get_variable(regularizer=layers.l2_regularizer(scale=1e-5),initializer=get_bilinear_filter([3,3,f_out,f_in],2),shape=[3,3,f_out,f_in],name='filter',trainable=is_training)\t\n\n\t\t\t\t# attach previous convolutional output to upsampling/deconvolutional output\n\t\t\t\ttmp = conv_out_list[-np.int32(i/nlayers_befPool)]\n\t\t\t\toutput_shape = tf.shape(tmp)\n\t\t\t\tnm_deconv_out = tf.nn.conv2d_transpose(nm_deconv_out,filter=W,output_shape=output_shape,strides=[1,2,2,1],padding='SAME')\n\t\t\t\tnm_deconv_out = layers.batch_norm(scope=scope,activation_fn=tf.nn.relu,inputs=nm_deconv_out,decay=0.9, center=True, scale=True, epsilon=1e-4, param_initializers={'beta_initializer':tf.zeros_initializer(),'gamma_initializer':tf.ones_initializer(),'moving_variance_initializer':tf.ones_initializer(),'moving_average_initializer':tf.zeros_initializer()}, param_regularizers={'beta_regularizer':None,'gamma_regularizer':layers.l2_regularizer(scale=1e-5)},is_training=is_training,trainable=is_training)\n\n\n\t\t\ttmp = layers.conv2d(tmp,num_outputs=f_out,kernel_size=[3,3],stride=[1,1],padding='SAME',normalizer_fn=layers.batch_norm, normalizer_params=batch_norm_params,weights_initializer=tf.random_normal_initializer(mean=0,stddev=np.sqrt(2/9/f_in)),weights_regularizer=layers.l2_regularizer(scale=1e-5),biases_initializer=None,scope=scope,trainable=is_training)\n\t\t\tnm_deconv_out = tmp + nm_deconv_out\n\n\n\t\telif i==deconv_layers:\n\t\t\tnm_deconv_out = layers.conv2d(nm_deconv_out,num_outputs=f_out,kernel_size=[3,3],stride=[1,1],padding='SAME',normalizer_fn=None,activation_fn=None,weights_initializer=tf.random_normal_initializer(mean=0,stddev=np.sqrt(2/9/f_in)),weights_regularizer=layers.l2_regularizer(scale=1e-5),biases_initializer=None,scope=scope,trainable=is_training)\n\n\n\t\telse:\n\t\t\tnm_deconv_out = layers.conv2d(nm_deconv_out,num_outputs=f_out,kernel_size=[3,3],stride=[1,1],padding='SAME',normalizer_fn=layers.batch_norm, normalizer_params=batch_norm_params,weights_initializer=tf.random_normal_initializer(mean=0,stddev=np.sqrt(2/9/f_in)),weights_regularizer=layers.l2_regularizer(scale=1e-5),biases_initializer=None,scope=scope,trainable=is_training)\n\n\n\n\treturn am_deconv_out, nm_deconv_out\n\n \n\ndef get_bilinear_filter(filter_shape, upscale_factor):\n    ##filter_shape is [width, height, num_in_channels, num_out_channels]\n    kernel_size = filter_shape[1]\n    ### Centre location of the filter for which value is calculated\n    if kernel_size % 2 == 1:\n        centre_location = upscale_factor - 1\n    else:\n        centre_location = upscale_factor - 0.5\n\n    x,y = np.meshgrid(np.arange(kernel_size),np.arange(kernel_size))\n    bilinear = (1 - abs((x - centre_location)/ upscale_factor)) * (1 - abs((y - centre_location)/ upscale_factor))\n    weights = np.tile(bilinear[:,:,None,None],(1,1,filter_shape[2],filter_shape[3]))\n\n    return tf.constant_initializer(weights)\n\n\n\n\n\n\n\n\n\n\n"
  },
  {
    "path": "model/dataloader.py",
    "content": "import pickle as pk\nimport os\nimport numpy as np\nimport tensorflow as tf\nimport skimage.transform as imgTform\nimport glob\nfrom scipy import io\n\n\ndef megaDepth_dataPipeline(num_subbatch_input, dir):\n\t# import ipdb; ipdb.set_trace()\n\t# locate all scenes \n\tdata_scenes1 = np.array(sorted(glob.glob(os.path.join(dir, '*'))))\n\n\t# scan scenes\n\t# sort scenes by number of training images in each\n\tscenes_size1 = np.array([len(os.listdir(i)) for i in data_scenes1])\n\tscenes_sorted1 = np.argsort(scenes_size1)\n\n\t# define scenes for training and testing\n\ttrain_scenes = data_scenes1[scenes_sorted1]\n\n\n\t# load data from each scene\n\t# locate each data minibatch in each sorted sc\n\ttrain_scenes_items = [sorted(glob.glob(os.path.join(sc, '*.pk'))) for sc in train_scenes]\n\ttrain_scenes_items = np.concatenate(train_scenes_items, axis=0)\n\n\ttrain_items = train_scenes_items\n\n\t### contruct training data pipeline\n\t# remove residual data over number of data in one epoch\n\tres_train_items = len(train_items) - (len(train_items) % num_subbatch_input)\n\ttrain_items = train_items[:res_train_items]\n\ttrain_data = md_construct_inputPipeline(train_items, flag_shuffle=True, batch_size=num_subbatch_input)\n\n\t# define re-initialisable iterator\n\titerator = tf.data.Iterator.from_structure(train_data.output_types, train_data.output_shapes)\n\tnext_element = iterator.get_next()\n\n\t# define initialisation for each iterator\n\ttrainData_init_op = iterator.make_initializer(train_data)\n\n\treturn next_element, trainData_init_op, len(train_items)\n\n\ndef _read_pk_function(filename):\n\twith open(filename, 'rb') as f:\n\t\tbatch_data = pk.load(f)\n\tinput = np.float32(batch_data['input'])\n\tdm = batch_data['dm']\n\tnm = np.float32(batch_data['nm'])\n\tcam = np.float32(batch_data['cam'])\n\tscaleX= batch_data['scaleX']\n\tscaleY = batch_data['scaleY']\n\tmask = np.float32(batch_data['mask'])\n\n\treturn input, dm, nm, cam, scaleX, scaleY, mask\n\ndef md_read_func(filename):\n\n\tinput, dm, nm, cam, scaleX, scaleY, mask = tf.py_func(_read_pk_function, [filename], [tf.float32, tf.float32, tf.float32, tf.float32, tf.float32, tf.float32, tf.float32])\t\n\n\tinput = tf.data.Dataset.from_tensor_slices(input[None])\n\tdm = tf.data.Dataset.from_tensor_slices(dm[None])\n\tnm = tf.data.Dataset.from_tensor_slices(nm[None])\n\tcam = tf.data.Dataset.from_tensor_slices(cam[None])\n\tscaleX = tf.data.Dataset.from_tensor_slices(scaleX[None])\n\tscaleY = tf.data.Dataset.from_tensor_slices(scaleY[None])\n\tmask = tf.data.Dataset.from_tensor_slices(mask[None])\n\n\treturn tf.data.Dataset.zip((input, dm, nm, cam, scaleX, scaleY, mask))\n\n\ndef md_preprocess_func(input, dm, nm, cam, scaleX, scaleY, mask):\n\n\tinput = input/255.\n\n\tnm = nm/127\n\n\treturn input, dm, nm, cam, scaleX, scaleY, mask\n\n\ndef md_construct_inputPipeline(items, batch_size, flag_shuffle=True):\n\tdata = tf.data.Dataset.from_tensor_slices(items)\n\tif flag_shuffle:\n\t\tdata = data.apply(tf.contrib.data.shuffle_and_repeat(buffer_size=100000))\n\telse:\n\t\tdata = data.repeat()\n\tdata = data.apply(tf.contrib.data.parallel_interleave(md_read_func, cycle_length=batch_size, block_length=1, sloppy=False ))\n\tdata = data.map(md_preprocess_func, num_parallel_calls=8 )\n\tdata = data.batch(batch_size).prefetch(4)\n\n\treturn data\n\n\n"
  },
  {
    "path": "model/lambSH_layer.py",
    "content": "import tensorflow as tf\n\n\n# am is the albedo map, which has shape (batch, height, width, 3[rgb]) \n# nm is the sparse normal map, which has shape (batch, height, width, 3[x,y,z])\n# L_SHcoeff contains the SH coefficients for environment illumination, using 2nd order SH. L_SHcoeff has shape (batch, 9, 3[rgb])\ndef lambSH_layer(am, nm, L_SHcoeffs, gamma):\n\n\t\"\"\" \n\ti = albedo * irradiance\n\tthe multiplication is elementwise\n\talbedo is given\n\tirraidance = n.T * M * n, where n is (x,y,z,1)\n\tM is contructed from some precomputed constants and L_SHcoeff, where M contains information about illuminations, clamped cosine and SH basis\n\t\"\"\"\n\n\t# M is only related with lighting\n\tc1 = tf.constant(0.429043,dtype=tf.float32)\n\tc2 = tf.constant(0.511664,dtype=tf.float32)\n\tc3 = tf.constant(0.743125,dtype=tf.float32)\n\tc4 = tf.constant(0.886227,dtype=tf.float32)\n\tc5 = tf.constant(0.247708,dtype=tf.float32)\n\n\t# each row have shape (batch, 4, 3)\n\tM_row1 = tf.stack([c1*L_SHcoeffs[:,8,:], c1*L_SHcoeffs[:,4,:], c1*L_SHcoeffs[:,7,:], c2*L_SHcoeffs[:,3,:]],axis=1)\n\tM_row2 = tf.stack([c1*L_SHcoeffs[:,4,:], -c1*L_SHcoeffs[:,8,:], c1*L_SHcoeffs[:,5,:], c2*L_SHcoeffs[:,1,:]],axis=1)\n\tM_row3 = tf.stack([c1*L_SHcoeffs[:,7,:], c1*L_SHcoeffs[:,5,:], c3*L_SHcoeffs[:,6,:], c2*L_SHcoeffs[:,2,:]],axis=1)\n\tM_row4 = tf.stack([c2*L_SHcoeffs[:,3,:], c2*L_SHcoeffs[:,1,:], c2*L_SHcoeffs[:,2,:], c4*L_SHcoeffs[:,0,:]-c5*L_SHcoeffs[:,6,:]],axis=1)\n\n\t# M is a 5d tensot with shape (batch,4,4,3[rgb]), the axis 1 and 2 are transposely equivalent\n\tM = tf.stack([M_row1,M_row2,M_row3,M_row4], axis=1)\n\n\t# find batch-spatial three dimensional mask of defined normals over nm\n\t# mask = tf.logical_not(tf.is_nan(nm[:,:,:,0]))\n\tmask = tf.not_equal(tf.reduce_sum(nm,axis=-1),0)\n\n\n\t# extend Cartesian to homogeneous coords and extend its last for rgb individual multiplication dimension, nm_homo have shape (total_npix, 4)\n\ttotal_npix = tf.shape(nm)[:3]\n\tones = tf.ones(total_npix)\n\tnm_homo = tf.concat([nm,tf.expand_dims(ones,axis=-1)], axis=-1)\n\n\t# contruct batch-wise flatten M corresponding with nm_homo, such that multiplication between them is batch-wise\n\tM = tf.expand_dims(tf.expand_dims(M,axis=1),axis=1)\n\n\n\t# expand M for broadcasting, such that M has shape (npix,4,4,3)\n\t# expand nm_homo, such that nm_homo has shape (npix,4,1,1)\n\tnm_homo = tf.expand_dims(tf.expand_dims(nm_homo,axis=-1),axis=-1)\n\t# tmp have shape (npix, 4, 3[rgb])\n\ttmp = tf.reduce_sum(nm_homo*M,axis=-3)\n\t# E has shape (npix, 3[rbg])\n\tE = tf.reduce_sum(tmp*nm_homo[:,:,:,:,0,:],axis=-2)\n\n\n\t# compute intensity by product between irradiance and albedo\n\ti = E*am\n\n\t# gamma correction\n\ti = tf.clip_by_value(i, 0., 1.) + tf.constant(1e-4)\n\ti = tf.pow(i,1./gamma)\n\n\treturn i, mask\n\n\n\n\n\n\n\n\n\n\n\n\n\n"
  },
  {
    "path": "model/loss_layer.py",
    "content": "# formulate loss function based on supplied ground truth and outputs from network\n\nimport importlib\nimport tensorflow as tf\nimport numpy as np\nimport os\nfrom model import SfMNet, lambSH_layer, pred_illuDecomp_layer, sup_illuDecomp_layer, reproj_layer\n\ndef loss_formulate(albedos, nm_pred, am_sup, nm_gt, inputs, dms, cams, scale_xs, scale_ys, masks, pair_label, preTrain_flag, am_smt_w_var, reproj_w_var, reg_loss_flag=True):\n\n\t# define gamma nonlinear mapping factor\n\tgamma = tf.constant(2.2)\n\n\talbedos = tf.nn.sigmoid(albedos) * masks + tf.constant(1e-4)\n\n\t### pre-process nm_pred such that in range (-1,1)\n\tnm_pred_norm = tf.sqrt(tf.reduce_sum(nm_pred**2, axis=-1, keepdims=True)+tf.constant(1.))\n\tnm_pred_xy = nm_pred / nm_pred_norm\n\tnm_pred_z = tf.constant(1.) / nm_pred_norm\n\tnm_pred_xyz = tf.concat([nm_pred_xy, nm_pred_z], axis=-1) * masks\n\n\t# selete normal map used in rendering - gt or pred\n\tnormals = nm_gt if preTrain_flag else nm_pred_xyz\n\n\n\t# reconstruct SH lightings from predicted statistical SH lighting model\n\tlighting_model = '../hdr_illu_pca'\n\tlighting_vectors = tf.constant(np.load(os.path.join(lighting_model,'pcaVector.npy')),dtype=tf.float32)\n\tlighting_means = tf.constant(np.load(os.path.join(lighting_model,'mean.npy')),dtype=tf.float32)\n\tlightings_var = tf.constant(np.load(os.path.join(lighting_model,'pcaVariance.npy')),dtype=tf.float32)\n\t\n\tif preTrain_flag:\n\t\tlightings = sup_illuDecomp_layer.illuDecomp(inputs,albedos,nm_gt,gamma)\n\telse:\n\t\tlightings =pred_illuDecomp_layer.illuDecomp(inputs,albedos,nm_pred_xyz,gamma,masks)\n\n\tlightings_pca = tf.matmul((lightings - lighting_means), pinv(lighting_vectors))\n\n\t# recompute lightings from lightins_pca which could add weak constraint on lighting reconstruction \n\tlightings = tf.matmul(lightings_pca,lighting_vectors) + lighting_means \n\n\t# reshape 27-D lightings to 9*3 lightings\n\tlightings = tf.reshape(lightings,[tf.shape(lightings)[0],9,3])\n\n\n\t### lighting prior loss\n\tvar = tf.reduce_mean(lightings_pca**2,axis=0)\n\n\tillu_prior_loss = tf.losses.absolute_difference(var, lightings_var)\n\n\tillu_prior_loss = tf.log(illu_prior_loss + 1.)\n\n\n\t### stereo supervision based on albedos reprojection consistancy\n\treproj_tb = tf.to_float(tf.equal(pair_label,tf.transpose(pair_label)))\n\treproj_tb = tf.cast(tf.matrix_set_diag(reproj_tb, tf.zeros([tf.shape(inputs)[0]])),tf.bool)\n\treproj_list = tf.where(reproj_tb)\n\timg1_inds = tf.expand_dims(reproj_list[:,0],axis=-1)\n\timg2_inds = tf.expand_dims(reproj_list[:,1],axis=-1)\n\talbedo1 = tf.gather_nd(albedos,img1_inds)\n\tdms1 = tf.gather_nd(dms,img1_inds)\n\tcams1 = tf.gather_nd(cams,img1_inds)\n\talbedo2 = tf.gather_nd(albedos,img2_inds)\n\tcams2 = tf.gather_nd(cams,img2_inds)\n\tscale_xs1 = tf.gather_nd(scale_xs, img1_inds)\n\tscale_xs2 = tf.gather_nd(scale_xs, img2_inds)\n\tscale_ys1 = tf.gather_nd(scale_ys, img1_inds)\n\tscale_ys2 = tf.gather_nd(scale_ys, img2_inds)\n\n\tinput1 = tf.gather_nd(inputs, img1_inds)\n\n\t# mask_indices contains indices for image index inside batch and spatial locations, and ignores the rgb channel index\n\treproj_albedo1, reproj_mask = reproj_layer.map_reproj(dms1,albedo2,cams1,cams2,scale_xs1,scale_xs2,scale_ys1,scale_ys2)\n\n\treproj_albedo1 = reproj_albedo1+tf.constant(1e-4) # numerical stable constant\n\n\n\n\t### scale intensities for each image\n\tnum_imgs = tf.shape(reproj_mask)[0]\n\tim_ = tf.constant(0)\n\toutput = tf.TensorArray(dtype=tf.float32,size=num_imgs)\t\n\n\tdef body(im_, output):\n\t\treproj_mask_ = reproj_mask[im_]\n\t\talbedo1_ = tf.boolean_mask(albedo1[im_],reproj_mask_)\n\t\treproj_albedo1_ = tf.boolean_mask(reproj_albedo1[im_],reproj_mask_)\n\n\n\t\tk = tf.reduce_sum(albedo1_*reproj_albedo1_,keepdims=True)/(tf.reduce_sum(reproj_albedo1_**2,keepdims=True)+tf.constant(1e-4))\n\n\t\toutput = output.write(im_,k)\n\t\tim_ += tf.constant(1)\n\n\t\treturn im_, output\n\n\tdef condition(im_, output):\n\t\treturn tf.less(im_,num_imgs)\n\n\t_,output = tf.while_loop(condition, body, loop_vars=[im_, output])\n\n\n\tks = tf.expand_dims(output.stack(), axis=-1)\n\n\n\n\talbedo1_pixels = tf.boolean_mask(albedo1, reproj_mask)\n\treproj_albedo1_pixels = tf.boolean_mask(reproj_albedo1*ks, reproj_mask)\n\treproj_err = tf.losses.mean_squared_error(cvtLab(albedo1_pixels), cvtLab(reproj_albedo1_pixels))\n\n\n\t### formulate loss based on paired batches ###\n\t# self-supervision based on intensity reconstruction\n\tshadings, renderings_mask = lambSH_layer.lambSH_layer(tf.ones_like(albedos), normals, lightings, 1.)\n\n\t# compare rendering intensity by Lab\n\tinputs_pixels = cvtLab(tf.boolean_mask(inputs,renderings_mask))\n\trenderings = cvtLab(tf.boolean_mask(tf.pow(albedos*shadings,1./gamma),renderings_mask))\n\trender_err = tf.losses.mean_squared_error(inputs_pixels,renderings)\n\n\n\t### compute rendering loss from cross-projected alebdo map\n\tcross_shadings = tf.gather_nd(shadings, img1_inds)\n\tinputs_pixels = cvtLab(tf.boolean_mask(input1,reproj_mask))\n\tcross_renderings = cvtLab(tf.boolean_mask(tf.pow(tf.nn.relu(cross_shadings*reproj_albedo1*ks), 1./gamma),reproj_mask))\n\tcross_render_err = tf.losses.mean_squared_error(inputs_pixels,cross_renderings)\n\n\n\t### measure smoothness of albedo map\n\tGx = tf.constant(1/2)*tf.expand_dims(tf.expand_dims(tf.constant([[-1,1]], dtype=tf.float32), axis=-1), axis=-1)\n\tGy = tf.constant(1/2)*tf.expand_dims(tf.expand_dims(tf.constant([[-1],[1]], dtype=tf.float32), axis=-1), axis=-1)\n\tGx_3 = tf.tile(Gx, multiples=(1,1,3,1))\t\n\tGy_3 = tf.tile(Gy, multiples=(1,1,3,1))\t\n\talbedo_lab = tf.reshape(cvtLab(tf.reshape(albedos,[-1,3])),[-1,200,200,3])\n\n\taGx = tf.nn.conv2d(albedos, Gx_3, padding='SAME', strides=(1,1,1,1))\n\taGy = tf.nn.conv2d(albedos, Gy_3, padding='SAME', strides=(1,1,1,1))\n\taGxy = tf.concat([aGx,aGy], axis=-1)\n\n\n\t# compute pixel-wise smoothness weights by angle distance between neighbour pixels' chromaticities\n\tinputs_pad = tf.pad(inputs, paddings=tf.constant([[0,0], [0,1], [0,1], [0,0]]))\n\tchroma_pad = tf.nn.l2_normalize(inputs_pad, axis=-1)\n\n\tchroma = chroma_pad[:,:-1,:-1,:]\n\tchroma_X = chroma_pad[:,:-1,1:,:]\n\tchroma_Y = chroma_pad[:,1:,:-1,:]\n\tchroma_Gx = tf.reduce_sum(chroma*chroma_X, axis=-1, keepdims=True)**tf.constant(2.) - tf.constant(1.)\n\tchroma_Gy = tf.reduce_sum(chroma*chroma_Y, axis=-1, keepdims=True)**tf.constant(2.) - tf.constant(1.)\n\tchroma_Gx = tf.exp(chroma_Gx / tf.constant(0.0001))\n\tchroma_Gy = tf.exp(chroma_Gy / tf.constant(0.0001))\n\tchroma_Gxy = tf.concat([chroma_Gx, chroma_Gy], axis=-1)\n\n\tint_pad = tf.reduce_sum(inputs_pad**tf.constant(2.), axis=-1, keepdims=True)\n\tint = int_pad[:,:-1,:-1,:]\n\tint_X = int_pad[:,:-1,1:,:]\n\tint_Y = int_pad[:,1:,:-1,:]\n\n\tint_Gx = tf.where(condition=int < int_X, x=int, y=int_X)\n\tint_Gy = tf.where(condition=int < int_Y, x=int, y=int_Y)\n\tint_Gx = tf.constant(1.) + tf.exp(- int_Gx / tf.constant(.8))\n\tint_Gy = tf.constant(1.) + tf.exp(- int_Gy / tf.constant(.8))\n\tint_Gxy = tf.concat([int_Gx, int_Gy], axis=-1)\n\n\tGxy_weights = int_Gxy * chroma_Gxy\n\talbedo_smt_error = tf.reduce_mean(tf.abs(aGxy)*Gxy_weights)\n\n\n\t### albedo map pseudo-supervision loss\n\tif preTrain_flag:\n\t\tam_loss = tf.constant(0.)\n\telse:\n\t\tamSup_mask = tf.not_equal(tf.reduce_sum(nm_gt,axis=-1),0)\n\t\tam_sup_pixel = cvtLab(tf.boolean_mask(am_sup, amSup_mask))\n\t\talbedos_pixel = cvtLab(tf.boolean_mask(albedos, amSup_mask))\n\t\tam_loss = tf.losses.mean_squared_error(am_sup_pixel, albedos_pixel)\n\n\n\n\t### regualarisation loss\n\treg_loss = sum(tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES))\n\n\n\t### compute nm_pred error\n\tnmSup_mask = tf.not_equal(tf.reduce_sum(nm_gt,axis=-1),0)\n\tnm_gt_pixel = tf.boolean_mask(nm_gt, nmSup_mask)\n\tnm_pred_pixel = tf.boolean_mask(nm_pred_xyz, nmSup_mask)\n\tnm_prod = tf.reduce_sum(nm_pred_pixel * nm_gt_pixel, axis=-1, keepdims=True)\t\n\tnm_cosValue = tf.constant(0.9999)\n\tnm_prod = tf.clip_by_value(nm_prod, -nm_cosValue, nm_cosValue)\n\tnm_angle = tf.acos(nm_prod) + tf.constant(1e-4)\n\tnm_loss = tf.reduce_mean(nm_angle**2)\n\n\n\n\t### compute gradient loss\n\tnm_pred_Gx = conv2d_nosum(nm_pred_xyz, Gx)\n\tnm_pred_Gy = conv2d_nosum(nm_pred_xyz, Gy)\n\tnm_pred_Gxy = tf.concat([nm_pred_Gx, nm_pred_Gy], axis=-1)\n\tnormals_Gx = conv2d_nosum(nm_gt, Gx)\n\tnormals_Gy = conv2d_nosum(nm_gt, Gy)\n\tnormals_Gxy = tf.concat([normals_Gx, normals_Gy], axis=-1)\n\n\tnm_pred_smt_error = tf.losses.mean_squared_error(nm_pred_Gxy, normals_Gxy)\n\n\n\t### total loss\n\trender_err *= tf.constant(.1)\n\treproj_err *= tf.constant(.05) * reproj_w_var\n\tcross_render_err *= tf.constant(.1)\n\tam_loss *= tf.constant(.1)\n\tillu_prior_loss *= tf.constant(.01)\n\talbedo_smt_error *= tf.constant(50.) * am_smt_w_var\n\tnm_pred_smt_error *= tf.constant(1.)\n\tnm_loss *= tf.constant(1.)\n\n\n\n\tif reg_loss_flag == True:\n\t\tloss = render_err + reproj_err + cross_render_err + reg_loss + illu_prior_loss + albedo_smt_error + nm_pred_smt_error + nm_loss + am_loss\n\telse:\n\t\tloss = render_err + reproj_err + cross_render_err + illu_prior_loss + albedo_smt_error + nm_pred_smt_error + nm_loss + am_loss\n\n\treturn lightings, albedos, nm_pred_xyz, loss, render_err, reproj_err, cross_render_err, reg_loss, illu_prior_loss, albedo_smt_error, nm_pred_smt_error, nm_loss, am_loss\n\n\n\n# input RGB is 2d tensor with shape (n_pix, 3)\ndef cvtLab(RGB):\n\n\t# threshold definition\n\tT = tf.constant(0.008856)\n\n\t# matrix for converting RGB to LUV color space\n\tcvt_XYZ = tf.constant([[0.412453,0.35758,0.180423],[0.212671,0.71516,0.072169],[0.019334,0.119193,0.950227]])\n\n\t# convert RGB to XYZ\n\tXYZ = tf.matmul(RGB,tf.transpose(cvt_XYZ))\n\n\t# normalise for D65 white point\n\tXYZ /= tf.constant([[0.950456, 1., 1.088754]])*100\n\n\tmask = tf.to_float(tf.greater(XYZ,T))\n\n\tfXYZ = XYZ**(1/3)*mask + (1.-mask)*(tf.constant(7.787)*XYZ + tf.constant(0.137931))\n\n\tM_cvtLab = tf.constant([[0., 116., 0.], [500., -500., 0.], [0., 200., -200.]])\n\n\tLab = tf.matmul(fXYZ, tf.transpose(M_cvtLab)) + tf.constant([[-16., 0., 0.]])\n\tmask = tf.to_float(tf.equal(Lab, tf.constant(0.)))\n\n\tLab += mask * tf.constant(1e-4)\n\n\treturn Lab\n\n\n\n\n\n# compute pseudo inverse for input matrix\ndef pinv(A, reltol=1e-6):\n\t# compute SVD of input A\n\ts, u, v = tf.svd(A)\n\n\t# invert s and clear entries lower than reltol*s_max\n\tatol = tf.reduce_max(s) * reltol\n\ts = tf.where(s>atol, s, atol*tf.ones_like(s))\n\ts_inv = tf.diag(1./s)\n\n\t# compute v * s_inv * u_t as psuedo inverse\n\treturn tf.matmul(v, tf.matmul(s_inv, tf.transpose(u)))\n\n\n\n# compute regular 2d convolution on 3d data\ndef conv2d_nosum(input, kernel):\n\tinput_x = input[:,:,:,0:1]\n\tinput_y = input[:,:,:,1:2]\n\tinput_z = input[:,:,:,2:3]\n\n\toutput_x = tf.nn.conv2d(input_x, kernel, strides=(1,1,1,1), padding='SAME')\n\toutput_y = tf.nn.conv2d(input_y, kernel, strides=(1,1,1,1), padding='SAME')\n\toutput_z = tf.nn.conv2d(input_z, kernel, strides=(1,1,1,1), padding='SAME')\n\n\treturn tf.concat([output_x,output_y,output_z], axis=-1)\n\n\n\n# compute regular 2d convolution on 3d data\ndef conv2d_nosum_2ch(input, kernel):\n\tinput_x = input[:,:,:,0:1]\n\tinput_y = input[:,:,:,1:2]\n\n\toutput_x = tf.nn.conv2d(input_x, kernel, strides=(1,1,1,1), padding='SAME')\n\toutput_y = tf.nn.conv2d(input_y, kernel, strides=(1,1,1,1), padding='SAME')\n\n\treturn tf.concat([output_x,output_y], axis=-1)\n\n\n\n\n\n\n\n\n"
  },
  {
    "path": "model/pred_illuDecomp_layer.py",
    "content": "import tensorflow as tf\n\n\n# am is the albedo map, which has shape (batch, height, width, 3[rgb]) \n# nm is the sparse normal map, which has shape (batch, height, width, 3[x,y,z])\n# L_SHcoeff contains the SH coefficients for environment illumination, using 2nd order SH. L_SHcoeff has shape (batch, 9, 3[rgb])\ndef illuDecomp(input, am, nm, gamma, masks):\n\n\t\"\"\" \n\ti = albedo * irradiance\n\tthe multiplication is elementwise\n\talbedo is given\n\tirraidance = n.T * M * n, where n is (x,y,z,1)\n\tM is contructed from some precomputed constants and L_SHcoeff, where M contains information about illuminations, clamped cosine and SH basis\n\t\"\"\"\n\n\t# compute shading by dividing input by albedo\n\tshadings = tf.pow(input,gamma)/am\n\t# perform clamping on resulted shading to guarantee its numerical range\n\tshadings = (tf.clip_by_value(shadings, 0., 1.) + tf.constant(1e-4)) * masks\n\n\n\t# compute shading by linear equation regarding nm and L_SHcoeffs\n\t# E(n) = c1*L22*(x**2-y**2) + (c3*z**2 - c5)*L20 + c4*L00 + 2*c1*L2-2*x*y + 2*c1*L21*x*z + 2*c1*L2-1*y*z + 2*c2*L11*x + 2*c2*L1-1*y + 2*c2*L10*z\n\t# E(n) = c4*L00 + 2*c2*y*L1-1 + 2*c2*z*L10 + 2*c2*x*L11 + 2*c1*x*y*L2-2 + 2*c1*y*z*L2-1 + (c3*z**2 - c5)*L20 + 2*c1*x*z*L21 + c1*(x**2-y**2)*L22 \n\tc1 = tf.constant(0.429043,dtype=tf.float32)\n\tc2 = tf.constant(0.511664,dtype=tf.float32)\n\tc3 = tf.constant(0.743125,dtype=tf.float32)\n\tc4 = tf.constant(0.886227,dtype=tf.float32)\n\tc5 = tf.constant(0.247708,dtype=tf.float32)\n\n\n\t# find defined pixels\n\tnum_iter = tf.shape(nm)[0]\n\toutput = tf.TensorArray(dtype=tf.float32, size=num_iter)\n\ti = tf.constant(0)\n\n\tdef condition(i, output):\n\t\treturn i<num_iter\n\n\tdef body(i, output):\n\t\tshadings_ = shadings[i]\n\t\tnm_ = nm[i]\n\t\tshadings_pixel = tf.reshape(shadings_, (-1,3))\n\t\tnm_ = tf.reshape(nm_, (-1,3))\n\n\t\t# E(n) = A*L_SHcoeffs\n\t\ttotal_npix = tf.shape(nm_)[0:1]\n\t\tones = tf.ones(total_npix)\n\t\tA = tf.stack([c4*ones, 2*c2*nm_[:,1], 2*c2*nm_[:,2], 2*c2*nm_[:,0], 2*c1*nm_[:,0]*nm_[:,1], 2*c1*nm_[:,1]*nm_[:,2], c3*nm_[:,2]**2-c5, 2*c1*nm_[:,2]*nm_[:,0], c1*(nm_[:,0]**2-nm_[:,1]**2)], axis=-1)\n\t\toutput = output.write(i, tf.matmul(pinv(A), shadings_pixel))\n\t\ti += tf.constant(1)\n\n\t\treturn i, output\n\n\t_, output = tf.while_loop(condition, body, loop_vars=[i,output])\n\tL_SHcoeffs = output.stack()\n\n\n\treturn tf.reshape(L_SHcoeffs, [-1,27])\n\n\n\ndef pinv(A, reltol=1e-6):\n\t# compute SVD of input A\n\ts, u, v = tf.svd(A)\n\n\t# invert s and clear entries lower than reltol*s_max\n\tatol = tf.reduce_max(s) * reltol\n\ts = tf.boolean_mask(s, s>atol)\n\ts_inv = tf.diag(1./s)\n\n\t# compute v * s_inv * u_t as psuedo inverse\n\treturn tf.matmul(v, tf.matmul(s_inv, tf.transpose(u)))\n\n\n\n\n"
  },
  {
    "path": "model/reproj_layer.py",
    "content": "# apply error mask in albedo reprojection\n\n\n# no rotation involved\n\n\n#### directly output flatten reprojected pixels and the reconstruction mask\n\n# the differentiable layer performing reprojection\n\nimport tensorflow as tf\nimport numpy as np\n\n# pc is n-by-3 matrix containing point could three locations\n# cam is the new camera parameters, whose f and p_a have shape (batch) and c has shape (batch, 2)\n# dm1 is the depth map associated with cam1 that is camera for output image, which has shape (batch, height, width)\n# img2 is the input image that acts as source image for reprojection, which has shape (batch, height, width, 3)\ndef map_reproj(dm1,map2,cam1,cam2,scale_x1,scale_x2,scale_y1,scale_y2):\n\tbatch_size = tf.shape(dm1)[0]\n\n\t# read camera parameters\n\tc1 = cam1[:,2:4]\n\tf1 = cam1[:,0]\n\tp_a1 = cam1[:,1] # ratio is width divided by height\n\tR1 = tf.reshape(cam1[:,4:13],[-1,3,3])\n\tt1 = cam1[:,13:]\n\n\tc2 = cam2[:,2:4]\n\tf2 = cam2[:,0]\n\tp_a2 = cam2[:,1]\n\tR2 = tf.reshape(cam2[:,4:13],[-1,3,3])\n\tt2 = cam2[:,13:]\n\n\t# project pixel points back to camera coords\n\t# u is the height and v is the width\n\t# u and v are scalars\n\tu1 = tf.shape(dm1)[1]\n\tv1 = tf.shape(dm1)[2]\n\n\t# convert u1 and v1 to float, convenient for computation\n\tu1 = tf.to_float(u1)\n\tv1 = tf.to_float(v1)\n\n\t### regular grid in output image\n\t# x increase towards right, y increase toward down\n\tvm,um = tf.meshgrid(tf.range(1.,v1+1.), tf.range(1.,u1+1.))\n\n\n\t# apply scaling factors on f\n\t# f1 = f1/(scale_x1+scale_y1)*2\n\t# f1 = tf.stack([f1, f1*p_a1],axis=-1)\n\tf1 = tf.stack([f1/scale_x1, f1/scale_y1*p_a1],axis=-1)\n\n\t# expand f1 (batch,2,1,1), to be consistant with dm \n\tf1 = tf.expand_dims(tf.expand_dims(f1,axis=-1),axis=-1)\n\t# expand c1 dimension (batch,2,1,1)\n\tc1 = tf.expand_dims(tf.expand_dims(c1,axis=-1),axis=-1)\n\t# expand vm and um to have shape (1,height,width)\n\tvm = tf.expand_dims(vm,axis=0)\t\n\tum = tf.expand_dims(um,axis=0)\t\n\n\t# compute 3D point x and y coordinates\n\t# Xm and Ym have shape (batch, height, width)\n\tXm = (vm-c1[:,0])/f1[:,0]*dm1\n\tYm = (um-c1[:,1])/f1[:,1]*dm1\n\n\t# the point cloud is (batch, 3, npix) matrix, each row is XYZ cam coords for one point\n\tpc = tf.stack([tf.contrib.layers.flatten(Xm), tf.contrib.layers.flatten(Ym), tf.contrib.layers.flatten(dm1)], axis=1)\n\n\t### transfer pc from coords of cam1 to cam2\n\t# construct homogeneous point cloud with shape batch-4-by-num_pix\n\tnum_pix = tf.shape(pc)[-1]\n\thomo_pc_c1 = tf.concat([pc, tf.ones((batch_size,1,num_pix), dtype=tf.float32)], axis=1)\n\n\t# both transformation matrix have shape batch-by-4-by-4, valid for multiplication with defined homogeneous point cloud\n\tlast_row = tf.tile(tf.constant([[[0,0,0,1]]],dtype=tf.float32), multiples=[batch_size,1,1])\n\tW_C_R_t1 = tf.concat([R1,tf.expand_dims(t1,axis=2)],axis=2)\n\tW_C_trans1 = tf.concat([W_C_R_t1, last_row], axis=1)\n\tW_C_R_t2 = tf.concat([R2,tf.expand_dims(t2,axis=2)],axis=2)\n\tW_C_trans2 = tf.concat([W_C_R_t2, last_row], axis=1)\n\n\t# batch dot product, output has shape (batch, 4, npix)\t\n\thomo_pc_c2 = tf.matmul(W_C_trans2, tf.matmul(tf.matrix_inverse(W_C_trans1), homo_pc_c1))\n\n\t### project point cloud to cam2 pixel coordinates\n\t# u in vertical and v in horizontal\n\tu2 = tf.shape(map2)[1]\n\tv2 = tf.shape(map2)[2]\n\n\t# convert u2 and v2 to float\n\tu2 = tf.to_float(u2)\n\tv2 = tf.to_float(v2)\n\t\n\t# f2 = f2/(scale_x2+scale_y2)*2\n\t# f2 = tf.stack([f2, f2*p_a2],axis=-1)\n\tf2 = tf.stack([f2/scale_x2, f2/scale_y2*p_a2],axis=-1)\n\n\t# construct intrics matrics, which has shape (batch, 3, 4)\n\tzeros = tf.zeros_like(f2[:,0],dtype=tf.float32)\t\n\tones = tf.ones_like(f2[:,0],tf.float32)\n\tk2 = tf.stack([tf.stack([f2[:,0],zeros,c2[:,0],zeros],axis=1), tf.stack([zeros,f2[:,1],c2[:,1],zeros],axis=1), tf.stack([zeros,zeros,ones,zeros],axis=1)],axis=1)\n\n\t## manual batch dot product\n\tk2 = tf.expand_dims(k2,axis=-1)\n\thomo_pc_c2 = tf.expand_dims(homo_pc_c2,axis=1)\n\t# homo_uv2 has shape (batch, 3, npix)\n\thomo_uv2 = tf.reduce_sum(k2*homo_pc_c2,axis=2)\n\n\t# the reprojected locations of regular grid in output image\n\t# both have shape (batch, npix)\n\tv_reproj = homo_uv2[:,0,:]/homo_uv2[:,2,:]\n\tu_reproj = homo_uv2[:,1,:]/homo_uv2[:,2,:]\n\n\t# u and v are flatten vector containing reprojected pixel locations\n\t# the u and v on same index compose one pixel \n\tu_valid = tf.logical_and(tf.logical_and(tf.logical_not(tf.is_nan(u_reproj)), u_reproj>0), u_reproj<u2-1)\n\tv_valid = tf.logical_and(tf.logical_and(tf.logical_not(tf.is_nan(v_reproj)), v_reproj>0), v_reproj<v2-1)\n\t# pixels has shape (batch, npix), indicating available reprojected pixels\n\tpixels = tf.logical_and(u_valid,v_valid)\n\n\t# pixels is bool indicator over original regular grid\n\t# v_reproj and u_reproj is x and y coordinates in source image\n\t# pixels, v_reproj and u_reproj are corresponded with each other by their indices\n\n\t### interpolation function based on source image img2\n\t# it has shape (total_npix, 3), the second dimension contains [img_inds, x, y]; we need to use img_inds to distinguish each pixel's request image\n\t# img_inds is 2d matrix with shape (batch, npix), containing img_ind for each (x,y) location\n\timg_inds = tf.tile(tf.expand_dims(tf.to_float(tf.range(batch_size)), axis=1), multiples=[1,num_pix])\n\trequest_points1 = tf.stack([tf.boolean_mask(img_inds,pixels), tf.boolean_mask(v_reproj,pixels), tf.boolean_mask(u_reproj,pixels)], axis=1)\n\n\n\n\t# the output is stacked flatten pixel values for channels\n\tre_proj_pixs = interpImg(request_points1, map2)\t\n\n\t# reconstruct original shaped re-projection map\n\tndims = tf.shape(map2)[3]\n\tshape = [batch_size, tf.to_int32(u1), tf.to_int32(v1),3]\n\n\tpixels = tf.reshape(pixels,shape=tf.stack([batch_size, tf.to_int32(u1), tf.to_int32(v1)],axis=0))\n\tindices = tf.to_int32(tf.where(tf.equal(pixels,True)))\n\n\tre_proj_pixs = tf.scatter_nd(updates=re_proj_pixs, indices=indices, shape=shape)\n\n\t# re_proj_pix is flatten reprojection results with shape (total_npix, 3)\n\t# indices contains first three indices in original image shape for each pixel in re_proj_pixs\n\treturn re_proj_pixs, pixels\n\n\n\ndef interpImg(unknown,data):\n\t# interpolate unknown data on pixel locations defined in unknown from known data with location defined in on regular grid\n\n\t# find neighbour pixels on regular grid\n\t# x is horizontal, y is vertical\n\timg_inds = tf.to_int32(unknown[:,0])\n\tx = unknown[:,1]\n\ty = unknown[:,2]\n\t# rgb_inds = tf.to_int32(unknown[:,3])\n\n\tlow_x = tf.to_int32(tf.floor(x))\n\thigh_x = tf.to_int32(tf.ceil(x))\n\tlow_y = tf.to_int32(tf.floor(y))\n\thigh_y = tf.to_int32(tf.ceil(y))\n\n\t# measure the weights for neighbourhood average based on distance\n\tdist_low_x = tf.expand_dims(x - tf.to_float(low_x), axis=-1)\n\tdist_high_x = tf.expand_dims(tf.to_float(high_x) - x, axis=-1)\n\tdist_low_y = tf.expand_dims(y - tf.to_float(low_y), axis=-1)\n\tdist_high_y = tf.expand_dims(tf.to_float(high_y) - y, axis=-1)\n\n\t# compute horizontal avarage\n\tavg_low_y = dist_low_x*tf.gather_nd(data, indices=tf.stack([img_inds,low_y,low_x],axis=1)) + dist_high_x*tf.gather_nd(data, indices=tf.stack([img_inds,low_y,high_x],axis=1))\n\tavg_high_y = dist_low_x*tf.gather_nd(data, indices=tf.stack([img_inds,high_y,low_x],axis=1)) + dist_high_x*tf.gather_nd(data, indices=tf.stack([img_inds,high_y,high_x],axis=1))\n\n\t# compute vertical average\n\tavg = dist_low_y*avg_low_y + dist_high_y*avg_high_y\n\n\treturn avg\n\n\n\n\n\n\n"
  },
  {
    "path": "model/sup_illuDecomp_layer.py",
    "content": "import tensorflow as tf\n\n\n# am is the albedo map, which has shape (batch, height, width, 3[rgb]) \n# nm is the sparse normal map, which has shape (batch, height, width, 3[x,y,z])\n# L_SHcoeff contains the SH coefficients for environment illumination, using 2nd order SH. L_SHcoeff has shape (batch, 9, 3[rgb])\ndef illuDecomp(input, am, nm, gamma):\n\n\t\"\"\" \n\ti = albedo * irradiance\n\tthe multiplication is elementwise\n\talbedo is given\n\tirraidance = n.T * M * n, where n is (x,y,z,1)\n\tM is contructed from some precomputed constants and L_SHcoeff, where M contains information about illuminations, clamped cosine and SH basis\n\t\"\"\"\n\n\t# compute shading by dividing input by albedo\n\tshadings = tf.pow(input,gamma)/(am)\n\t# perform clamping on resulted shading to guarantee its numerical range\n\tshadings = tf.clip_by_value(shadings, 0., 1.) + tf.constant(1e-4)\n\n\n\t# compute shading by linear equation regarding nm and L_SHcoeffs\n\tc1 = tf.constant(0.429043,dtype=tf.float32)\n\tc2 = tf.constant(0.511664,dtype=tf.float32)\n\tc3 = tf.constant(0.743125,dtype=tf.float32)\n\tc4 = tf.constant(0.886227,dtype=tf.float32)\n\tc5 = tf.constant(0.247708,dtype=tf.float32)\n\n\n\t# find defined pixels\n\tmask = tf.not_equal(tf.reduce_sum(nm,axis=-1),0)\n\tnum_iter = tf.shape(mask)[0]\n\toutput = tf.TensorArray(dtype=tf.float32, size=num_iter)\n\ti = tf.constant(0)\n\n\tdef condition(i, output):\n\t\treturn i<num_iter\n\n\tdef body(i, output):\n\t\tmask_ = mask[i]\n\t\tshadings_ = shadings[i]\n\t\tnm_ = nm[i]\n\t\tshadings_pixel = tf.boolean_mask(shadings_, mask_)\n\t\tnm_ = tf.boolean_mask(nm_, mask_)\n\n\t\t# E(n) = A*L_SHcoeffs\n\t\ttotal_npix = tf.shape(nm_)[0:1]\n\t\tones = tf.ones(total_npix)\n\t\tA = tf.stack([c4*ones, 2*c2*nm_[:,1], 2*c2*nm_[:,2], 2*c2*nm_[:,0], 2*c1*nm_[:,0]*nm_[:,1], 2*c1*nm_[:,1]*nm_[:,2], c3*nm_[:,2]**2-c5, 2*c1*nm_[:,2]*nm_[:,0], c1*(nm_[:,0]**2-nm_[:,1]**2)], axis=-1)\n\t\toutput = output.write(i, tf.matmul(pinv(A), shadings_pixel))\n\t\ti += tf.constant(1)\n\n\t\treturn i, output\n\n\t_, output = tf.while_loop(condition, body, loop_vars=[i,output])\n\tL_SHcoeffs = output.stack()\n\n\treturn tf.reshape(L_SHcoeffs, [-1,27])\n\n\n\ndef pinv(A, reltol=1e-6):\n\t# compute SVD of input A\n\ts, u, v = tf.svd(A)\n\n\t# invert s and clear entries lower than reltol*s_max\n\tatol = tf.reduce_max(s) * reltol\n\ts = tf.boolean_mask(s, s>atol)\n\ts_inv = tf.diag(1./s)\n\n\t# compute v * s_inv * u_t as psuedo inverse\n\treturn tf.matmul(v, tf.matmul(s_inv, tf.transpose(u)))\n\n\n\n"
  },
  {
    "path": "pre_train_model/.keep",
    "content": ""
  },
  {
    "path": "test_demo.py",
    "content": "import os\nimport numpy as np\nimport tensorflow as tf\nimport cv2\nfrom skimage import io\nimport argparse\nfrom model import SfMNet, lambSH_layer, pred_illuDecomp_layer\nfrom utils import render_sphere_nm\n\n\nparser = argparse.ArgumentParser(description='InverseRenderNet')\nparser.add_argument('--image', help='Path to test image')\nparser.add_argument('--mask', help='Path to image mask')\nparser.add_argument('--model', help='Path to trained model')\nparser.add_argument('--output', help='Folder saving outputs')\n\n\nargs = parser.parse_args()\n\nimg_path = args.image\nmask_path = args.mask\n\nimg = io.imread(img_path)\nmask = io.imread(mask_path)\n\n\ndst_dir = args.output\nos.makedirs(dst_dir)\n\ninput_height = 200\ninput_width = 200\nori_height, ori_width = img.shape[:2]\n\nif ori_height / ori_width >1:\n    scale = ori_width / 200\n    input_height = np.int32(scale * 200)\nelse:\n    scale = ori_height / 200\n    input_width = np.int32(scale * 200)\n\n\n# compute pseudo inverse for input matrix\ndef pinv(A, reltol=1e-6):\n\t# compute SVD of input A\n\ts, u, v = tf.svd(A)\n\n\t# invert s and clear entries lower than reltol*s_max\n\tatol = tf.reduce_max(s) * reltol\n\ts = tf.boolean_mask(s, s>atol)\n\ts_inv = tf.diag(1./s)\n\n\t# compute v * s_inv * u_t as psuedo inverse\n\treturn tf.matmul(v, tf.matmul(s_inv, tf.transpose(u)))\n\n\nimport ipdb; ipdb.set_trace()\ninputs_var = tf.placeholder(tf.float32, (None, input_height, input_width, 3))\nmasks_var = tf.placeholder(tf.float32, (None, input_height, input_width, 1))\nam_deconvOut, nm_deconvOut = SfMNet.SfMNet(inputs=inputs_var,is_training=False, height=input_height, width=input_width, n_layers=30, n_pools=4, depth_base=32)\n\n\n# separate albedo, error mask and shadow mask from deconvolutional output\nalbedos = am_deconvOut\nnm_pred = nm_deconvOut\n\ngamma = tf.constant(2.2)\n\n# post-process on raw albedo and nm_pred\nalbedos = tf.nn.sigmoid(albedos) * masks_var + tf.constant(1e-4)\n\nnm_pred_norm = tf.sqrt(tf.reduce_sum(nm_pred**2, axis=-1, keepdims=True)+tf.constant(1.))\nnm_pred_xy = nm_pred / nm_pred_norm\nnm_pred_z = tf.constant(1.) / nm_pred_norm\nnm_pred_xyz = tf.concat([nm_pred_xy, nm_pred_z], axis=-1) * masks_var\n\n\n# compute illumination\nlighting_model = 'illu_pca'\nlighting_vectors = tf.constant(np.load(os.path.join(lighting_model,'pcaVector.npy')),dtype=tf.float32)\nlighting_means = tf.constant(np.load(os.path.join(lighting_model,'mean.npy')),dtype=tf.float32)\t\nlightings = pred_illuDecomp_layer.illuDecomp(inputs_var, albedos, nm_pred_xyz, gamma, masks_var)\n\n\nlightings_pca = tf.matmul((lightings - lighting_means), pinv(lighting_vectors))\nlightings = tf.matmul(lightings_pca,lighting_vectors) + lighting_means \n# reshape 27-D lightings to 9*3 lightings\nlightings = tf.reshape(lightings,[tf.shape(lightings)[0],9,3])\n\n# visualisations\nshading, _ = lambSH_layer.lambSH_layer(tf.ones_like(albedos), nm_pred_xyz, lightings, 1.)\nnm_sphere = tf.constant(render_sphere_nm.render_sphere_nm(100,1),dtype=tf.float32)\nnm_sphere = tf.tile(nm_sphere, (tf.shape(inputs_var)[0],1,1,1))\nlighting_recon, _ = lambSH_layer.lambSH_layer(tf.ones_like(nm_sphere), nm_sphere, lightings, 1.)\n\n\nirn_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='conv') + tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='am') + tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='nm')\nmodel_path = tf.train.get_checkpoint_state(args.model).model_checkpoint_path\n\ntotal_loss = 0 \nsess = tf.InteractiveSession()\nsaver = tf.train.Saver(irn_vars)\nsaver.restore(sess, model_path)\n\n\n# evaluation\nori_img = img\nori_height, ori_width = ori_img.shape[:2]\nimg = cv2.resize(img, (input_width, input_height))\nimg = np.float32(img)/255.\nimg = img[None, :, :, :]\nmask = cv2.resize(mask, (input_width, input_height), cv2.INTER_NEAREST)\nmask = np.float32(mask==255)[None,:,:,None]\n\n[albedos_val, nm_pred_val, lighting_recon_val, shading_val] = sess.run([albedos, nm_pred_xyz, lighting_recon, shading], feed_dict={inputs_var:img, masks_var:mask})\n\n\n# post-process results\nnm_pred_val = (nm_pred_val+1.)/2.\n\nalbedos_val = cv2.resize(albedos_val[0], (ori_width, ori_height))\nshading_val = cv2.resize(shading_val[0], (ori_width, ori_height))\nlighting_recon_val = lighting_recon_val[0]\nnm_pred_val = cv2.resize(nm_pred_val[0], (ori_width, ori_height))\n\n\nalbedos_val = (albedos_val-albedos_val.min()) / (albedos_val.max()-albedos_val.min())\n\nalbedos_val = np.uint8(albedos_val*255.)\nshading_val = np.uint8(shading_val*255.)\nlighting_recon_val = np.uint8(lighting_recon_val*255.)\nnm_pred_val = np.uint8(nm_pred_val*255.)\n\ninput_path = os.path.join(dst_dir, 'img.png')\nio.imsave(input_path, ori_img)\nalbedo_path = os.path.join(dst_dir, 'albedo.png')\nio.imsave(albedo_path, albedos_val)\nshading_path = os.path.join(dst_dir, 'shading.png')\nio.imsave(shading_path, shading_val)\nnm_pred_path = os.path.join(dst_dir, 'nm_pred.png')\nio.imsave(nm_pred_path, nm_pred_val)\nlighting_path = os.path.join(dst_dir, 'lighting.png')\nio.imsave(lighting_path, lighting_recon_val)\n\n\n"
  },
  {
    "path": "test_iiw.py",
    "content": "import json\nimport os\nimport numpy as np\nimport tensorflow as tf\nimport importlib\nimport cv2\nfrom skimage import io\nimport argparse\nfrom model import SfMNet, lambSH_layer, pred_illuDecomp_layer\nfrom glob import glob\nfrom utils.whdr import compute_whdr\n\n\nparser = argparse.ArgumentParser(description='InverseRenderNet')\nparser.add_argument('--iiw', help='Root directory for iiw-dataset')\nparser.add_argument('--model', help='Path to trained model')\n\n\nargs = parser.parse_args()\n\niiw = args.iiw\ntest_ids = np.load('iiw_test_ids.npy')\n\n\n\ninput_height = 200\ninput_width = 200\n\n\n\n# compute pseudo inverse for input matrix\ndef pinv(A, reltol=1e-6):\n\t# compute SVD of input A\n\ts, u, v = tf.svd(A)\n\n\t# invert s and clear entries lower than reltol*s_max\n\tatol = tf.reduce_max(s) * reltol\n\ts = tf.boolean_mask(s, s>atol)\n\ts_inv = tf.diag(1./s)\n\n\t# compute v * s_inv * u_t as psuedo inverse\n\treturn tf.matmul(v, tf.matmul(s_inv, tf.transpose(u)))\n\n\n\ninputs_var = tf.placeholder(tf.float32, (None, input_height, input_width, 3))\nmasks_var = tf.placeholder(tf.float32, (None, input_height, input_width, 1))\ntrain_flag = tf.placeholder(tf.bool, ())\nam_deconvOut, _ = SfMNet.SfMNet(inputs=inputs_var,is_training=train_flag, height=input_height, width=input_width, n_layers=30, n_pools=4, depth_base=32)\n\n\n# separate albedo, error mask and shadow mask from deconvolutional output\nalbedos = am_deconvOut\n\n# post-process on raw albedo and nm_pred\nalbedos = tf.nn.sigmoid(albedos) * masks_var + tf.constant(1e-4)\n\nirn_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='conv') + tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='am') + tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='nm')\nmodel_path = tf.train.get_checkpoint_state(args.model).model_checkpoint_path\n\ntotal_loss = 0 \nsess = tf.InteractiveSession()\nsaver = tf.train.Saver(irn_vars)\nsaver.restore(sess, model_path)\n\n\nfor counter, test_id in enumerate(test_ids):\n    img_file = str(test_id)+'.png'\n    judgement_file = str(test_id)+'.json'\n\n    img_path = os.path.join(iiw, 'data', img_file)\n    judgement_path = os.path.join(iiw, 'data', judgement_file)\n\n    img = io.imread(img_path)\n    judgement = json.load(open(judgement_path))\n\n    ori_width, ori_height = img.shape[:2]\n\n    img = cv2.resize(img, (input_width, input_height))\n    img = np.float32(img)/255.\n    img = img[None, :, :, :]\n    mask = np.ones((1, input_height, input_width, 1), np.bool)\n\n\n    [albedos_val] = sess.run([albedos], feed_dict={train_flag:False, inputs_var:img, masks_var:mask})\n\n    albedos_val = cv2.resize(albedos_val[0], (ori_width, ori_height))\n\n    albedos_val = (albedos_val-albedos_val.min()) / (albedos_val.max()-albedos_val.min())\n    albedos_val = albedos_val/2+.5\n\n\n    loss = compute_whdr(albedos_val, judgement)\n    total_loss += loss\n    print('whdr:{:f}\\twhdr_avg:{:f}'.format(loss, total_loss/(counter+1)))\n\n\nprint(\"IIW TEST WHDR %f\"%(total_loss/len(test_ids)))\n\n\n"
  },
  {
    "path": "train.py",
    "content": "# also predict shadow mask and error mask\n\n# no rotation\n\n\n#### compute albedo reproj loss only on reprojection available area; compute reconstruction and its loss only based on defined area\n\n\nimport tensorflow as tf\nimport importlib\nimport os\nimport pickle as pk\nimport sys\nimport numpy as np\nimport time\nimport argparse\nfrom PIL import Image\nimport glob\nfrom model import SfMNet, lambSH_layer, pred_illuDecomp_layer, loss_layer, dataloader\n\n\nparser = argparse.ArgumentParser(description='InverseRenderNet')\nparser.add_argument('--n_batch', '-n', help='number of minibatch', type=int)\nparser.add_argument('--data_path', '-p', help='Path to training data')\nparser.add_argument('--train_mode', '-m', help='specify the phase for training (pre-train/self-train)', choices={'pre-train', 'self-train'})\n\n\nargs = parser.parse_args()\n\ndef main():\n\n\tinputs_shape = (5,200,200,3)\n\n\tnext_element, trainData_init_op, num_train_batches = dataloader.megaDepth_dataPipeline(args.n_batch, args.data_path)\n\n\tinputs_var = tf.reshape(next_element[0], (-1, inputs_shape[1], inputs_shape[2], inputs_shape[3]))\n\tdms_var = tf.reshape(next_element[1], (-1, inputs_shape[1], inputs_shape[2]))\n\tnms_var = tf.reshape(next_element[2], (-1, inputs_shape[1], inputs_shape[2], 3))\n\tcams_var = tf.reshape(next_element[3], (-1, 16))\n\tscaleXs_var = tf.reshape(next_element[4], (-1,))\n\tscaleYs_var = tf.reshape(next_element[5], (-1,))\n\tmasks_var = tf.reshape(next_element[6], (-1, inputs_shape[1], inputs_shape[2]))\n\n\t# var helping cross projection\n\tpair_label_var = tf.constant(np.repeat(np.arange(args.n_batch),inputs_shape[0])[:,None], dtype=tf.float32)\n\t# weights for smooth loss and am_consistency loss\n\tam_smt_w_var = tf.placeholder(tf.float32, ())\n\treproj_w_var = tf.placeholder(tf.float32, ())\n\n\t# mask out sky in inputs and nms\n\tmasks_var_4d = tf.expand_dims(masks_var, axis=-1)\n\tinputs_var *= masks_var_4d\n\tnms_var *= masks_var_4d\n\n\t# inverserendernet\n\tif args.train_mode == 'pre-train':\n\t\tam_deconvOut, nm_deconvOut = SfMNet.SfMNet(inputs=inputs_var,is_training=True, height=inputs_shape[1], width=inputs_shape[2], name='pre_train_IRN/', n_layers=30, n_pools=4, depth_base=32)\n\n\t\tam_sup = tf.zeros_like(am_deconvOut)\n\t\tpreTrain_flag = True\n\n\n\telif args.train_mode == 'self-train':\n\t\tam_deconvOut, nm_deconvOut = SfMNet.SfMNet(inputs=inputs_var,is_training=True, height=inputs_shape[1], width=inputs_shape[2], name='IRN/', n_layers=30, n_pools=4, depth_base=32)\n\n\t\tam_sup, _ = SfMNet.SfMNet(inputs=inputs_var,is_training=False, height=inputs_shape[1], width=inputs_shape[2], name='pre_train_IRN/', n_layers=30, n_pools=4, depth_base=32)\n\t\tam_sup = tf.nn.sigmoid(am_sup) * masks_var_4d + tf.constant(1e-4)\n\n\t\tpreTrain_flag = False\n\n\t# separate albedo, error mask and shadow mask from deconvolutional output\n\talbedoMaps = am_deconvOut[:,:,:,:3]\n\n\t# formulate loss\n\tlight_SHCs, albedoMaps, nm_preds, loss, render_err, reproj_err, cross_render_err, reg_loss, illu_prior_loss, albedo_smt_error, nm_smt_loss, nm_loss, am_loss = loss_layer.loss_formulate(albedoMaps, nm_deconvOut, am_sup, nms_var, inputs_var, dms_var, cams_var, scaleXs_var, scaleYs_var, masks_var_4d, pair_label_var, True, am_smt_w_var, reproj_w_var, reg_loss_flag=True)\n\n\n\t# defined traning loop\n\tepochs = 30\n\tnum_batches = num_train_batches\n\tnum_subbatch = args.n_batch\n\tnum_iters = np.int32(np.ceil(num_batches/num_subbatch))\n\n\n\t# training op\n\tglobal_step = tf.Variable(1,name='global_step',trainable=False)\n\n\ttrain_step = tf.contrib.layers.optimize_loss(loss, optimizer=tf.train.AdamOptimizer(learning_rate=.05, epsilon=1e-1), learning_rate=None, global_step=global_step)\n\n\t# define saver for saving and restoring\n\tirn_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='IRN') if args.train_mode == 'self-train' else tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='pre_train_IRN')\n\tsaver = tf.train.Saver(irn_vars)\n\n\t# define session\n\tconfig = tf.ConfigProto(allow_soft_placement=True)\n\tconfig.gpu_options.allow_growth = True\n\tsess = tf.InteractiveSession(config=config)\n\n\t# train from scratch or keep training trained model\n\ttf.local_variables_initializer().run()\n\ttf.global_variables_initializer().run()\n\n\tassignOps = []\n\tif args.train_mode == 'self-train':\n\t\t# load am_sup net\n\t\tpreTrain_irn_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='pre_train_IRN')\n\t\tsaver_loadOldVar = tf.train.Saver(preTrain_irn_vars)\n\t\tsaver_loadOldVar.restore(sess, 'pre_train_model/model.ckpt')\n\n\t\t# import ipdb; ipdb.set_trace()\n\t\t# duplicate pre_train model\n\t\twith tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):\n\t\t\tvars = tf.contrib.framework.list_variables('pre_train_model')\n\t\t\tfor var_name, _ in vars:\n\t\t\t\tvar = tf.contrib.framework.load_variable('pre_train_model', var_name)\n\t\t\t\tnew_var_name = var_name.replace('pre_train_IRN', 'IRN')\n\n\t\t\t\tnew_var = tf.get_variable(name=new_var_name)\n\t\t\t\tassignOps += [new_var.assign(var)]\n\n\t\tsess.run(assignOps)\n\n\n\n\n\t# start training\n\ttrainData_init_op.run()\n\tdst_dir = 'irn_model' if args.train_mode == 'self-train' else 'pre_train_model'\n\tfor i in range(1,epochs+1):\n\n\t\tloss_avg = 0\n\t\tf = open('cost.txt','a')\n\n\t\t# graduately update weights if pre-training\n\t\treproj_weight = .2 + np.clip(.8 * (i-16)/14, 0., .8) if args.train_mode == 'pre-train' else 1.\n\t\tam_smt_weight = .2 + np.clip(.8 * (i-1)/14, 0., .8) if args.train_mode == 'pre-train' else 1.\n\n\t\tfor j in range(1,num_iters+1):\n\t\t\tstart_time = time.time()\n\n\t\t\t# train\n\t\t\t[loss_val, reg_loss_val, render_err_val, reproj_err_val, cross_render_err_val, illu_prior_val, albedo_smt_error_val, nm_smt_loss_val, nm_loss_val, am_loss_val] = sess.run([train_step, reg_loss, render_err, reproj_err, cross_render_err, illu_prior_loss, albedo_smt_error, nm_smt_loss, nm_loss, am_loss], feed_dict={am_smt_w_var:am_smt_weight, reproj_w_var:reproj_weight})\n\t\t\tloss_avg += loss_val\n\n\t\t\t# log\n\t\t\tif j % 1 == 0:\n\t\t\t\tprint('iter %d/%d loop %d/%d took %.3fs' % (i,epochs,j,num_iters,time.time()-start_time))\n\t\t\t\tprint('\\tloss_avg = %f, loss = %f' % (loss_avg / j,loss_val))\n\t\t\t\tprint('\\t\\treg_loss = %f, render_err = %f, reproj_err = %f, cross_render_err = %f, illu_prior = %f, albedo_smt_error = %f, nm_smt_loss = %f, nm_loss = %f, am_loss = %f' % (reg_loss_val, render_err_val, reproj_err_val, cross_render_err_val, illu_prior_val, albedo_smt_error_val, nm_smt_loss_val, nm_loss_val, am_loss_val))\n\n\t\t\t\tf.write('iter %d/%d loop %d/%d took %.3fs\\n\\tloss_avg = %f, loss = %f\\n\\t\\treg_loss = %f, render_err = %f, reproj_err = %f, cross_render_err = %f, illu_prior = %f, albedo_smt_error = %f, nm_smt_loss = %f, nm_loss = %f, am_loss = %f\\n' % (i,epochs,j,num_iters,time.time()-start_time,loss_avg/j, loss_val, reg_loss_val, render_err_val, reproj_err_val, cross_render_err_val, illu_prior_val, albedo_smt_error_val, nm_smt_loss_val, nm_loss_val, am_loss_val))\n\n\t\tf.close()\n\n\t\t# save model every 10 iterations\n\t\tsaver.save(sess,os.path.join(dst_dir, 'model.ckpt'))\n\n\nif __name__ == '__main__':\n\tmain()\n\n\n\n\n\n"
  },
  {
    "path": "utils/render_sphere_nm.py",
    "content": "import numpy as np\n\ndef render_sphere_nm(radius, num):\n\t# nm is a batch of normal maps\n\tnm = []\n\n\tfor i in range(num):\n\t\t### hemisphere\n\t\theight = 2*radius\n\t\twidth = 2*radius\n\t\tcentre = radius\n\t\tx_grid, y_grid = np.meshgrid(np.arange(1.,2*radius+1), np.arange(1.,2*radius+1))\n\t\t# grids are (-radius, radius)\n\t\tx_grid -= centre\n\t\t# y_grid -= centre\n\t\ty_grid = centre - y_grid\n\t\t# scale range of h and w grid in (-1,1)\n\t\tx_grid /= radius\n\t\ty_grid /= radius\n\t\tdist = 1 - (x_grid**2+y_grid**2)\n\t\tmask = dist > 0\n\t\tz_grid = np.ones_like(mask) * np.nan\n\t\tz_grid[mask] = np.sqrt(dist[mask])\n\n\t\t# remove xs and ys by masking out nans in zs\n\t\tx_grid[~(mask)] = np.nan\n\t\ty_grid[~(mask)] = np.nan\n\n\t\t# concatenate normal map\n\t\tnm.append(np.stack([x_grid,y_grid,z_grid],axis=2))\n\n\n\n\t\t### sphere \n\t\t# span the regular grid for computing azimuth and zenith angular map\n\t\t# height = 2*radius\n\t\t# width = 2*radius\n\t\t# centre = radius\n\t\t# h_grid, v_grid = np.meshgrid(np.arange(1.,2*radius+1), np.arange(1.,2*radius+1))\n\t\t# # grids are (-radius, radius)\n\t\t# h_grid -= centre\n\t\t# # v_grid -= centre\n\t\t# v_grid = centre - v_grid\n\t\t# # scale range of h and v grid in (-1,1)\n\t\t# h_grid /= radius\n\t\t# v_grid /= radius\n\n\t\t# # z_grid is linearly spread along theta/zenith in range (0,pi)\n\t\t# dist_grid = np.sqrt(h_grid**2+v_grid**2)\n\t\t# dist_grid[dist_grid>1] = np.nan\n\t\t# theta_grid = dist_grid * np.pi\n\t\t# z_grid = np.cos(theta_grid)\n\n\t\t# rho_grid = np.arctan2(v_grid,h_grid)\n\t\t# x_grid = np.sin(theta_grid)*np.cos(rho_grid)\n\t\t# y_grid = np.sin(theta_grid)*np.sin(rho_grid)\n\n\t\t# # concatenate normal map\n\t\t# nm.append(np.stack([x_grid,y_grid,z_grid],axis=2))\n\n\n\t# construct batch\n\tnm = np.stack(nm,axis=0)\n\n\n\n\treturn nm\n\n"
  },
  {
    "path": "utils/whdr.py",
    "content": "#!/usr/bin/env python2.7\n#\n# This is an implementation of the WHDR metric proposed in this paper:\n#\n#     Sean Bell, Kavita Bala, Noah Snavely. \"Intrinsic Images in the Wild\". ACM\n#     Transactions on Graphics (SIGGRAPH 2014). http://intrinsic.cs.cornell.edu.\n#\n# Please cite the above paper if you find this code useful.  This code is\n# released under the MIT license (http://opensource.org/licenses/MIT).\n#\n\n\nimport sys\nimport json\nimport argparse\nimport numpy as np\nfrom PIL import Image\n\n\ndef compute_whdr(reflectance, judgements, delta=0.10):\n    \"\"\" Return the WHDR score for a reflectance image, evaluated against human\n    judgements.  The return value is in the range 0.0 to 1.0, or None if there\n    are no judgements for the image.  See section 3.5 of our paper for more\n    details.\n\n    :param reflectance: a numpy array containing the linear RGB\n    reflectance image.\n\n    :param judgements: a JSON object loaded from the Intrinsic Images in\n    the Wild dataset.\n\n    :param delta: the threshold where humans switch from saying \"about the\n    same\" to \"one point is darker.\"\n    \"\"\"\n\n    points = judgements['intrinsic_points']\n    comparisons = judgements['intrinsic_comparisons']\n    id_to_points = {p['id']: p for p in points}\n    rows, cols = reflectance.shape[0:2]\n\n    error_sum = 0.0\n    weight_sum = 0.0\n\n    for c in comparisons:\n        # \"darker\" is \"J_i\" in our paper\n        darker = c['darker']\n        if darker not in ('1', '2', 'E'):\n            continue\n\n        # \"darker_score\" is \"w_i\" in our paper\n        weight = c['darker_score']\n        if weight <= 0 or weight is None:\n            continue\n\n        point1 = id_to_points[c['point1']]\n        point2 = id_to_points[c['point2']]\n        if not point1['opaque'] or not point2['opaque']:\n            continue\n\n        # convert to grayscale and threshold\n        l1 = max(1e-10, np.mean(reflectance[\n            int(point1['y'] * rows), int(point1['x'] * cols), ...]))\n        l2 = max(1e-10, np.mean(reflectance[\n            int(point2['y'] * rows), int(point2['x'] * cols), ...]))\n\n        # convert algorithm value to the same units as human judgements\n        if l2 / l1 > 1.0 + delta:\n            alg_darker = '1'\n        elif l1 / l2 > 1.0 + delta:\n            alg_darker = '2'\n        else:\n            alg_darker = 'E'\n\n        if darker != alg_darker:\n            error_sum += weight\n        weight_sum += weight\n\n    if weight_sum:\n        return error_sum / weight_sum\n    else:\n        return None\n\n\ndef load_image(filename, is_srgb=True):\n    \"\"\" Load an image that is either linear or sRGB-encoded. \"\"\"\n\n    if not filename:\n        raise ValueError(\"Empty filename\")\n    image = np.asarray(Image.open(filename)).astype(np.float) / 255.0\n    if is_srgb:\n        return srgb_to_rgb(image)\n    else:\n        return image\n\n\ndef srgb_to_rgb(srgb):\n    \"\"\" Convert an sRGB image to a linear RGB image \"\"\"\n\n    ret = np.zeros_like(srgb)\n    idx0 = srgb <= 0.04045\n    idx1 = srgb > 0.04045\n    ret[idx0] = srgb[idx0] / 12.92\n    ret[idx1] = np.power((srgb[idx1] + 0.055) / 1.055, 2.4)\n    return ret\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser(\n        description=(\n            'Evaluate an intrinsic image decomposition using the WHDR metric presented in:\\n'\n            '    Sean Bell, Kavita Bala, Noah Snavely. \"Intrinsic Images in the Wild\".\\n'\n            '    ACM Transactions on Graphics (SIGGRAPH 2014).\\n'\n            '    http://intrinsic.cs.cornell.edu.\\n'\n            '\\n'\n            'The output is in the range 0.0 to 1.0.'\n        )\n    )\n\n    parser.add_argument(\n        'reflectance', metavar='<reflectance.png>',\n        help='reflectance image to be evaluated')\n\n    parser.add_argument(\n        'judgements', metavar='<judgements.json>',\n        help='human judgements JSON file')\n\n    parser.add_argument(\n        '-l', '--linear', action='store_true', required=False,\n        help='assume the reflectance image is linear, otherwise assume sRGB')\n\n    parser.add_argument(\n        '-d', '--delta', metavar='<float>', type=float, required=False, default=0.10,\n        help='delta threshold (default 0.10)')\n\n    if len(sys.argv) < 2:\n        parser.print_help()\n        sys.exit(1)\n\n    args = parser.parse_args()\n    reflectance = load_image(filename=args.reflectance, is_srgb=(not args.linear))\n    judgements = json.load(open(args.judgements))\n\n    whdr = compute_whdr(reflectance, judgements, args.delta)\n    print(whdr)\n"
  }
]