Estimator comparison (attached reparam.)

tizian · Jul 28, 2022 · 1cd7a4a · 1cd7a4a
1 parent e6296da
commit 1cd7a4a
Show file tree

Hide file tree

Showing 10 changed files with 803 additions and 0 deletions.
diff --git a/.gitignore b/.gitignore
@@ -33,3 +33,4 @@ __pycache__
 results/.ipynb_checkpoints
 results/estimator_comparison_simple
 results/estimator_comparison_optimization
+results/estimator_comparison_attached_reparam
diff --git a/results/Estimator comparison (attached reparameterization).ipynb b/results/Estimator comparison (attached reparameterization).ipynb
@@ -0,0 +1,168 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "53581a91",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-07-28T11:30:15.423479Z",
+     "start_time": "2022-07-28T11:30:15.091516Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import mitsuba as mi\n",
+    "mi.set_variant(\"llvm_ad_rgb\")\n",
+    "import drjit as dr\n",
+    "\n",
+    "import os\n",
+    "base_dir = 'estimator_comparison_attached_reparam'\n",
+    "if not os.path.exists(base_dir):\n",
+    "    os.makedirs(base_dir)\n",
+    "    \n",
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "import cmap_diff\n",
+    "%config InlineBackend.figure_formats = ['svg']\n",
+    "%matplotlib inline\n",
+    "    \n",
+    "mi.Thread.thread().logger().set_log_level(mi.LogLevel.Warn)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5737c51e",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-07-28T11:31:29.924896Z",
+     "start_time": "2022-07-28T11:30:15.431353Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "scene_names = [\n",
+    "    ('attached_disc_roughness', 'plane.bsdf.alpha.data'),\n",
+    "    ('attached_disc_normalmap', 'plane.bsdf.normalmap.data'),\n",
+    "]\n",
+    "\n",
+    "methods = [\n",
+    "    ('Detached BSDF sampling',           {'method': 'bs_detached'}),\n",
+    "    ('Attached BSDF sampling (naive)',   {'method': 'bs_attached'}),\n",
+    "    ('Attached BSDF sampling (reparam)', {'method': 'bs_attached_reparam',\n",
+    "                                          'reparam_kappa': 1e6, 'reparam_rays': 48}),\n",
+    "]\n",
+    "\n",
+    "for scene_name, param_key in scene_names:\n",
+    "    print(\"*\", scene_name)\n",
+    "    scene_dir = \"{}/{}\".format(base_dir, scene_name)\n",
+    "    if not os.path.exists(scene_dir):\n",
+    "        os.makedirs(scene_dir)\n",
+    "    \n",
+    "    \n",
+    "    scene = mi.load_file('scenes/{}.xml'.format(scene_name), res=256)\n",
+    "    params = mi.traverse(scene)\n",
+    "    params.keep([param_key])\n",
+    "    params[param_key] = dr.maximum(0.1, params[param_key])\n",
+    "    \n",
+    "    # Primal rendering\n",
+    "    integrator = mi.load_dict({'type': 'estimator_comparison', 'method': 'primal_mis', 'hide_emitters': True})\n",
+    "    image = mi.render(scene, params, integrator=integrator, seed=0, spp=32)\n",
+    "    mi.util.convert_to_bitmap(image, uint8_srgb=False).write('{}/primal.exr'.format(scene_dir))\n",
+    "\n",
+    "    # Compare various methods ...\n",
+    "    for method_name, method_dict in methods:\n",
+    "        print(\"   \", method_name)\n",
+    "        integrator_dict = {'type': 'estimator_comparison', 'hide_emitters': True}\n",
+    "        for k, v in method_dict.items():\n",
+    "            integrator_dict[k] = v\n",
+    "        integrator = mi.load_dict(integrator_dict)\n",
+    "\n",
+    "        # Differentiable rendering ...\n",
+    "        dr.enable_grad(params[param_key])\n",
+    "        image = mi.render(scene, params, integrator=integrator, seed=0, spp=256, antithetic_pass=False)\n",
+    "        # ... and propagate back to input parameters\n",
+    "        dr.backward(image)\n",
+    "        param_grad = dr.grad(params[param_key])\n",
+    "        dr.set_grad(params[param_key], 0.0)\n",
+    "        dr.disable_grad(params[param_key])\n",
+    "        \n",
+    "        mi.util.convert_to_bitmap(param_grad, uint8_srgb=False).write('{}/{}.exr'.format(scene_dir, method_dict['method']))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4cf0ff0c",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2022-07-28T11:27:06.540158Z",
+     "start_time": "2022-07-28T11:27:05.821131Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "scales = [0.5, 2]\n",
+    "\n",
+    "for scene_idx, scene_name in enumerate([k for k, _ in scene_names]):\n",
+    "    scene_dir = \"{}/{}\".format(base_dir, scene_name)\n",
+    "    \n",
+    "    fig, axes = plt.subplots(ncols=4, figsize=(10,3))\n",
+    "    for ax in axes:\n",
+    "        ax.set_xticks([]); ax.set_yticks([])\n",
+    "        \n",
+    "    image_primal = np.array(mi.Bitmap(\"{}/primal.exr\".format(scene_dir)))\n",
+    "    image_primal = np.clip(image_primal**(1/2.2), 0.0, 1.0) # Crude gamma correction\n",
+    "    axes[0].imshow(image_primal)\n",
+    "    axes[0].set_title(\"Primal\")\n",
+    "    \n",
+    "    for method_idx, method_name in enumerate([v['method'] for k, v in methods]):\n",
+    "        vminmax = scales[scene_idx]\n",
+    "        image_grad = np.array(mi.Bitmap(\"{}/{}.exr\".format(scene_dir, method_name)))\n",
+    "        data_ = image_grad[:,:,0] if len(image_grad.shape) == 3 else image_grad\n",
+    "        im = axes[method_idx+1].imshow(data_, cmap='diff', vmin=-vminmax, vmax=+vminmax)\n",
+    "    \n",
+    "    axes[1].set_title(\"Detached\")\n",
+    "    axes[2].set_title(\"Attached (naïve)\")\n",
+    "    axes[3].set_title(\"Attached (reparam.)\")\n",
+    "    plt.suptitle(scene_name, weight='bold', size=14)\n",
+    "    \n",
+    "    outname = '{}/comparison.jpg'.format(scene_dir)\n",
+    "    plt.savefig(outname, dpi=300, pad_inches=0.1, bbox_inches='tight')\n",
+    "        \n",
+    "    plt.tight_layout()\n",
+    "    plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c70f0559",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "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.9.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/results/scenes/attached_disc_normalmap.xml b/results/scenes/attached_disc_normalmap.xml
@@ -0,0 +1,57 @@
+<scene version="2.0.0">
+    <integrator type="path"/>
+
+    <sensor type="perspective">
+        <string name="fov_axis" value="smaller"/>
+        <float name="near_clip" value="0.001"/>
+        <float name="far_clip" value="1000"/>
+        <float name="focus_distance" value="1000"/>
+        <float name="fov" value="60"/>
+        <transform name="to_world">
+            <lookat origin="0, 0, 2"
+                    target="0, 0, 0"
+                    up    ="0, 1, 0"/>
+            <rotate x="1" angle="-70"/>
+        </transform>
+        <sampler type="independent">
+            <integer name="sample_count" value="32"/>
+        </sampler>
+        <film type="hdrfilm">
+            <integer name="width" value="512"/>
+            <integer name="height" value="512"/>
+            <rfilter type="gaussian"/>
+            <string name="pixel_format" value="rgb"/>
+        </film>
+    </sensor>
+
+    <shape type="obj" id="plane">
+        <string name="filename" value="meshes/rounded_plane.obj"/>
+
+        <boolean name="face_normals" value="true"/>
+
+        <bsdf type="normalmap" id="nmap">
+            <texture name="normalmap" type="bitmap">
+                <string name="filename" value="textures/normalmap_lily.png"/>
+                <boolean name="raw" value="true"/>
+            </texture>
+
+            <bsdf type="roughconductor">
+                <float name="alpha" value="0.08"/>
+                <boolean name="differential_sampling" value="true"/>
+                <boolean name="sample_visible" value="true"/>
+            </bsdf>
+        </bsdf>
+
+        <transform name="to_world">
+            <scale x="1" y="1" z="1"/>
+            <translate z="-0.2"/>
+        </transform>
+    </shape>
+
+    <emitter type="envmap">
+        <string name="filename" value="textures/envmap.exr"/>
+        <transform name="to_world">
+            <rotate y="1" angle="90"/>
+        </transform>
+    </emitter>
+</scene>
diff --git a/results/scenes/attached_disc_roughness.xml b/results/scenes/attached_disc_roughness.xml
@@ -0,0 +1,67 @@
+<scene version="2.0.0">
+
+    <integrator type="path"/>
+
+    <sensor type="perspective">
+        <string name="fov_axis" value="smaller"/>
+        <float name="near_clip" value="0.001"/>
+        <float name="far_clip" value="1000"/>
+        <float name="focus_distance" value="1000"/>
+        <float name="fov" value="60"/>
+        <transform name="to_world">
+            <lookat origin="0, 0, 2"
+                    target="0, 0, 0"
+                    up    ="0, 1, 0"/>
+            <rotate x="1" angle="-70"/>
+        </transform>
+        <sampler type="independent">
+            <integer name="sample_count" value="32"/>
+        </sampler>
+        <film type="hdrfilm">
+            <integer name="width" value="512"/>
+            <integer name="height" value="512"/>
+            <rfilter type="gaussian"/>
+            <string name="pixel_format" value="rgb"/>
+        </film>
+    </sensor>
+
+    <shape type="obj" id="plane">
+        <string name="filename" value="meshes/rounded_plane.obj"/>
+
+        <boolean name="face_normals" value="true"/>
+
+        <bsdf type="roughconductor">
+            <rgb name="specular_reflectance" value="0.8, 0.8, 0.3"/>
+
+            <texture name="alpha" type="bitmap">
+                <string name="filename" value="textures/textured_roughness.jpg"/>
+                <boolean name="raw" value="true"/>
+            </texture>
+
+            <boolean name="differential_sampling" value="true"/>
+            <boolean name="principled_roughness_mapping" value="true"/>
+            <boolean name="sample_visible" value="true"/>
+        </bsdf>
+
+        <transform name="to_world">
+            <scale x="1" y="1" z="1"/>
+            <translate z="-0.2"/>
+        </transform>
+    </shape>
+
+    <shape type="sphere">
+        <float name="radius" value="1.2"/>
+        <point name="center" value="0, 3, -1.5"/>
+
+        <bsdf type="diffuse">
+            <spectrum name="reflectance" value="0.01"/>
+        </bsdf>
+    </shape>
+
+    <emitter type="envmap">
+        <string name="filename" value="textures/envmap.exr"/>
+        <transform name="to_world">
+            <rotate y="1" angle="90"/>
+        </transform>
+    </emitter>
+</scene>
diff --git a/results/scenes/textures/normalmap_lily.png b/results/scenes/textures/normalmap_lily.png
diff --git a/results/scenes/textures/textured_roughness.jpg b/results/scenes/textures/textured_roughness.jpg
diff --git a/src/python/python/ad/__init__.py b/src/python/python/ad/__init__.py
@@ -1,3 +1,4 @@
 from .integrators import *
 from .optimizers import *
 from .reparam import *
+from .attached_reparam import *