eval.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.argv = ['eval.py', '--checkpoint_name', 'fleece', '--yarn_name', 'fleece']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Rendering NN\n",
      "Batch (0/8)\n",
      "Batch (1/8)\n",
      "Batch (2/8)\n",
      "Batch (3/8)\n",
      "Batch (4/8)\n",
      "Batch (5/8)\n",
      "Batch (6/8)\n",
      "Batch (7/8)\n",
      "Rendering Ref\n",
      "Batch (0/8)\n",
      "Batch (1/8)\n",
      "Batch (2/8)\n",
      "Batch (3/8)\n",
      "Batch (4/8)\n",
      "Batch (5/8)\n",
      "Batch (6/8)\n",
      "Batch (7/8)\n"
     ]
    }
   ],
   "source": [
    "from config import device, variant\n",
    "\n",
    "import mitsuba as mi\n",
    "import drjit as dr\n",
    "mi.set_variant(variant)\n",
    "\n",
    "import torch\n",
    "\n",
    "import numpy as np\n",
    "from mitsuba import ScalarTransform4f as sT\n",
    "\n",
    "from utils.geometry import create_single_yarn\n",
    "\n",
    "from config.parameters import get_fiber_parameters\n",
    "from bsdf.neuralyarn import NeuralYarn\n",
    "from network.model import Model_M, Model_T\n",
    "from network.wrapper import MiModelWrapper\n",
    "\n",
    "import torch\n",
    "\n",
    "# import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "import argparse\n",
    "import os\n",
    "\n",
    "parser = argparse.ArgumentParser(description=\"Fitting the RDM\")\n",
    "\n",
    "parser.add_argument(\"--checkpoint_name\", help=\"Checkpoint name to store outputs\")\n",
    "parser.add_argument(\"--yarn_name\", help=\"Name of the yarn defined in config/parameters.py\")\n",
    "parser.add_argument(\"--batch_size\", type=int, help='Samples per pixel per batch rendering', default=32)\n",
    "parser.add_argument(\"--num_batches\", type=int, help='Number of batches for rendering', default=8)\n",
    "\n",
    "args = parser.parse_args()\n",
    "\n",
    "output_dir = os.path.join('checkpoints/', args.checkpoint_name)\n",
    "parameters = get_fiber_parameters(args.yarn_name)\n",
    "\n",
    "model_m = Model_M().to(device)\n",
    "model_t = Model_T().to(device)\n",
    "\n",
    "model_m.load_state_dict(torch.load(os.path.join(output_dir, 'model_m.pth'), weights_only=True))\n",
    "model_t.load_state_dict(torch.load(os.path.join(output_dir, 'model_t.pth'), weights_only=True))\n",
    "\n",
    "npz = np.load(os.path.join(output_dir, 'pdf.npz'))\n",
    "kappa_R, beta_M, gamma_M, kappa_M = npz['kappa_R'], npz['beta_M'], npz['gamma_M'], npz['kappa_M']\n",
    "\n",
    "mlp_m = MiModelWrapper(model_m, activation=dr.exp)\n",
    "mlp_t = MiModelWrapper(model_t, activation=lambda x: 1.0 / (1.0 + dr.exp(-x)))\n",
    "\n",
    "neuralyarn = NeuralYarn.create(parameters, mlp_m, mlp_t, kappa_R, beta_M, gamma_M, kappa_M)\n",
    "\n",
    "# Disable megekernel mode\n",
    "def mega_kernel(state):\n",
    "    dr.set_flag(dr.JitFlag.LoopRecord, state)\n",
    "    dr.set_flag(dr.JitFlag.VCallRecord, state)\n",
    "    dr.set_flag(dr.JitFlag.VCallOptimize, state)\n",
    "mega_kernel(False)\n",
    "\n",
    "scene_dict = {\n",
    "    'type': 'scene',\n",
    "    'integrator': {\n",
    "        'type': 'path',\n",
    "        'max_depth': -1,\n",
    "        'rr_depth': 9999999,\n",
    "        'hide_emitters': True,\n",
    "    },\n",
    "    'light1': {\n",
    "        'type': 'sphere',\n",
    "        'center': [0, 10, 10],\n",
    "        'radius': 2.5,\n",
    "        'emitter': {\n",
    "            'type': 'area',\n",
    "            'radiance': {\n",
    "                'type': 'rgb',\n",
    "                'value': 30,\n",
    "            }\n",
    "        }\n",
    "    },\n",
    "    'light2': {\n",
    "        'type': 'sphere',\n",
    "        'center': [0, -10, -10],\n",
    "        'radius': 2.5,\n",
    "        'emitter': {\n",
    "            'type': 'area',\n",
    "            'radiance': {\n",
    "                'type': 'rgb',\n",
    "                'value': 30,\n",
    "            }\n",
    "        }\n",
    "    },\n",
    "    'sensor': {\n",
    "        'type': 'perspective',\n",
    "        'to_world': sT.look_at(\n",
    "            origin=[0, 0, 10],\n",
    "            target=[0, 0, 0],\n",
    "            up=[0, 1, 0]\n",
    "        ),\n",
    "        'film': {\n",
    "            'type': 'hdrfilm',\n",
    "            'width': 256,\n",
    "            'height': 256,\n",
    "        },\n",
    "    },\n",
    "    'yarn': {\n",
    "        'type':'linearcurve',\n",
    "        'filename': './curves/ply.txt',\n",
    "        'bsdf': neuralyarn\n",
    "    }\n",
    "}\n",
    "\n",
    "os.makedirs(os.path.join(output_dir, 'images'), exist_ok=True)\n",
    "\n",
    "scene = mi.load_dict(scene_dict)\n",
    "\n",
    "batch_sample = args.batch_size\n",
    "samples = args.num_batches\n",
    "\n",
    "pred_image = None\n",
    "with torch.no_grad():\n",
    "    with dr.suspend_grad():\n",
    "        for i in range(samples):\n",
    "            pred_image_ = mi.render(scene, spp=batch_sample, seed=i)\n",
    "            try:\n",
    "                pred_image += pred_image_\n",
    "            except:\n",
    "                print('Rendering NN')\n",
    "                pred_image = pred_image_\n",
    "            print(f'Batch ({i+1}/{samples})')\n",
    "        pred_image /= samples \n",
    "\n",
    "        mi.util.write_bitmap(os.path.join(output_dir, 'images/pred.png'), pred_image)\n",
    "\n",
    "\n",
    "scene_dict_ref = scene_dict.copy()\n",
    "scene_dict_ref['yarn'] = create_single_yarn(parameters)\n",
    "scene_ref = mi.load_dict(scene_dict_ref)\n",
    "\n",
    "true_image = None\n",
    "with torch.no_grad():\n",
    "    with dr.suspend_grad():\n",
    "        for i in range(samples):\n",
    "            true_image_ = mi.render(scene_ref, spp=batch_sample, seed=i)\n",
    "            try:\n",
    "                true_image += true_image_\n",
    "            except:\n",
    "                print('Rendering Ref')\n",
    "                true_image = true_image_\n",
    "            print(f'Batch ({i+1}/{samples})')\n",
    "        true_image /= samples \n",
    "\n",
    "        mi.util.write_bitmap(os.path.join(output_dir, 'images/true.png'), true_image)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "neuralyarn",
   "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}