diff --git a/.gitignore b/.gitignore index 14954ca8..f7f10826 100644 --- a/.gitignore +++ b/.gitignore @@ -14,6 +14,8 @@ archive *.ply eval out +__pycache__ +try-runs # evaluation: temp/ @@ -25,3 +27,6 @@ vids/ stylegan3/results eg3d/results eg3d_results + +*runs/ +*ipynb diff --git a/dataset_preprocessing/abo/.gitignore b/dataset_preprocessing/abo/.gitignore new file mode 100644 index 00000000..c4c4ffc6 --- /dev/null +++ b/dataset_preprocessing/abo/.gitignore @@ -0,0 +1 @@ +*.zip diff --git a/dataset_preprocessing/abo/preprocess_abo_cameras.py b/dataset_preprocessing/abo/preprocess_abo_cameras.py new file mode 100644 index 00000000..d2eaacf6 --- /dev/null +++ b/dataset_preprocessing/abo/preprocess_abo_cameras.py @@ -0,0 +1,130 @@ +# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +############################################################# + +# Usage: python dataset_preprocessing/shapenet/preprocess_cars_cameras.py --source ~/downloads/cars_train --dest /data/cars_preprocessed + +############################################################# + + +# from distutils.debug import DEBUG +import json +import numpy as np +import os +from tqdm import tqdm +import argparse +from ipdb import set_trace as st + +def list_recursive(folderpath): + return [os.path.join(folderpath, filename) for filename in os.listdir(folderpath)] + +DEBUG=True + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument("--source", type=str) + parser.add_argument("--max_images", type=int, default=None) + args = parser.parse_args() + + # Parse cameras + dataset_path = args.source + cameras = {} + + blender2opencv = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]]) + w, h = 512, 512 + + #### this is for directly listing + # for scene_folder_path in list_recursive(dataset_path): + # if not os.path.isdir(scene_folder_path): continue + + #### this is using predefined list, to avoid folders that the datageneration is not complete + all_data = [] + resolution = 512 + for split in ['train', 'val']: + with open(os.path.join(dataset_path, 'meta', f'abo_{resolution}_{split}.txt')) as f: + scans = [line.rstrip() for line in f.readlines()] + all_data += scans + print(len(all_data), all_data) + for scene_folder_path_rel in all_data: + scene_folder_path = os.path.join(dataset_path, scene_folder_path_rel) + # st() # the sibling folder with rgb should be mesh>: no, only intrinsics and pose + + pointcloud_csv = os.path.join(scene_folder_path,'sample', f"pc.csv") + assert os.path.isfile(pointcloud_csv) + pc_relative_path = os.path.relpath(pointcloud_csv, dataset_path) + # print(pc_relative_path) + + with open(os.path.join(scene_folder_path,'render', f"transforms.json"), 'r') as f: + meta = json.load(f) + # print(meta.keys()) ['camera_angle_x', 'frames'] + # print(meta ['frames'][0]['file_path']) + + focal = .5 * w / np.tan(0.5 * meta['camera_angle_x']) + # intrinsic_for_all = np.array([[focal, 0, w / 2], [0, focal, h / 2], [0, 0, 1]]) + ### IMPORTANT!!! EG3D use the intrinsics that agnostic of original resolution!! + # intrinsics = np.array( + # [[focal / orig_img_size, 0.00000000e+00, cx / orig_img_size], + # [0.00000000e+00, focal / orig_img_size, cy / orig_img_size], + # [0.00000000e+00, 0.00000000e+00, 1.00000000e+00]] + # ).tolist() + intrinsic_for_all = np.array( + [[focal / w, 0.00000000e+00, (w / 2)/w], + [0.00000000e+00, focal / h, (h / 2)/h], + [0.00000000e+00, 0.00000000e+00, 1.00000000e+00]] + ).tolist() + st() + + # continue + # for rgb_path in list_recursive(os.path.join(scene_folder_path, 'render')): + for frame in meta ['frames']: + rgb_path = frame['file_path'] + relative_path = os.path.relpath(rgb_path, dataset_path) + print(relative_path) + + intrinsics = intrinsic_for_all + pose = (np.array(frame['transform_matrix'])@blender2opencv).tolist() + # print(len(pose)) + + cameras[relative_path] = {'pose': pose, 'intrinsics': intrinsics, 'scene-name': os.path.basename(scene_folder_path),\ + 'pc_csv':pc_relative_path} + # if DEBUG: + # break + + with open(os.path.join(dataset_path, 'cameras.json'), 'w') as outfile: + json.dump(cameras, outfile, indent=4) + + + camera_dataset_file = os.path.join(args.source, 'cameras.json') + + with open(camera_dataset_file, "r") as f: + cameras = json.load(f) # same camera file as saved above + + dataset = {'labels':[]} + # max_images = args.max_images if args.max_images is not None else len(cameras) + max_images = len(cameras) + for i, filename in tqdm(enumerate(cameras), total=max_images): + if (max_images is not None and i >= max_images): break + + pose = np.array(cameras[filename]['pose']) + intrinsics = np.array(cameras[filename]['intrinsics']) + label = np.concatenate([pose.reshape(-1), intrinsics.reshape(-1)]).tolist() + + image_path = os.path.join(args.source, filename) + pc_rel_path = cameras[filename]['pc_csv'] + dataset["labels"].append([filename, label, pc_rel_path]) + # also append pointcloud filename, but need to check with the dataset class too + + # print(dataset) + # check cameras/dataset + + with open(os.path.join(args.source, 'dataset.json'), "w") as f: + json.dump(dataset, f, indent=4) + diff --git a/dataset_preprocessing/abo/read_data.ipynb b/dataset_preprocessing/abo/read_data.ipynb new file mode 100644 index 00000000..b6dcf680 --- /dev/null +++ b/dataset_preprocessing/abo/read_data.ipynb @@ -0,0 +1,3352 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import pandas as pd\n", + "import os \n", + "import json\n", + "import torch \n", + "# csv = '/home/xuyi/Data/renderer/output_blender/lego/sample/pc.csv'\n", + "# base_dir = '/home/xuyi/Data/renderer/output_blender/lego'\n", + "# csv = '/home/xuyi/Data/renderer/output_partnet/3769/sample/pc.csv'\n", + "# base_dir = '/home/xuyi/Data/renderer/output_partnet/3769'\n", + "# csv_f = '/home/xuyi/Data/renderer/output_abo/B01N6AQX0A/sample/pc.csv'\n", + "csv_f = '/home/xuyi/Data/renderer/output_abo/B01D3C7Z4A/sample/pc.csv'\n", + "base_dir = '/home/xuyi/Data/renderer/output_abo/B07DBJX741'" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([[[-1.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00],\n", + " [ 0.0000e+00, 7.7414e-01, -6.3302e-01, 1.0305e+00],\n", + " [ 0.0000e+00, -6.3302e-01, -7.7414e-01, 1.2602e+00],\n", + " [ 0.0000e+00, 0.0000e+00, 0.0000e+00, 1.0000e+00]],\n", + "\n", + " [[-9.9211e-01, -9.7025e-02, 7.9338e-02, -1.2915e-01],\n", + " [-1.2533e-01, 7.6803e-01, -6.2803e-01, 1.0224e+00],\n", + " [ 0.0000e+00, -6.3302e-01, -7.7414e-01, 1.2602e+00],\n", + " [ 0.0000e+00, 0.0000e+00, 0.0000e+00, 1.0000e+00]],\n", + "\n", + " [[-9.6858e-01, -1.9241e-01, 1.5756e-01, -2.5649e-01],\n", + " [-2.4869e-01, 7.4939e-01, -6.1365e-01, 9.9895e-01],\n", + " [-7.4506e-09, -6.3355e-01, -7.7370e-01, 1.2595e+00],\n", + " [ 0.0000e+00, 0.0000e+00, 0.0000e+00, 1.0000e+00]],\n", + "\n", + " ...,\n", + "\n", + " [[-9.2978e-01, 2.8240e-01, -2.3615e-01, 3.8442e-01],\n", + " [ 3.6812e-01, 7.1326e-01, -5.9645e-01, 9.7095e-01],\n", + " [ 0.0000e+00, -6.4149e-01, -7.6713e-01, 1.2488e+00],\n", + " [ 0.0000e+00, 0.0000e+00, 0.0000e+00, 1.0000e+00]],\n", + "\n", + " [[-9.6858e-01, 1.9154e-01, -1.5861e-01, 2.5820e-01],\n", + " [ 2.4869e-01, 7.4600e-01, -6.1777e-01, 1.0057e+00],\n", + " [ 0.0000e+00, -6.3781e-01, -7.7020e-01, 1.2538e+00],\n", + " [ 0.0000e+00, 0.0000e+00, 0.0000e+00, 1.0000e+00]],\n", + "\n", + " [[-9.9212e-01, 9.6802e-02, -7.9603e-02, 1.2958e-01],\n", + " [ 1.2533e-01, 7.6630e-01, -6.3015e-01, 1.0258e+00],\n", + " [ 3.7253e-09, -6.3515e-01, -7.7239e-01, 1.2574e+00],\n", + " [ 0.0000e+00, 0.0000e+00, 0.0000e+00, 1.0000e+00]]],\n", + " dtype=torch.float64)" + ] + }, + "execution_count": 12, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "## read transform.json\n", + "with open(os.path.join(base_dir,'render', f\"transforms.json\"), 'r') as f:\n", + " meta = json.load(f)\n", + "all_poses=[]\n", + "blender2opencv = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])\n", + "for frame in meta ['frames']:\n", + " # rgb_path = frame['file_path']\n", + " # relative_path = os.path.relpath(rgb_path, dataset_path)\n", + " # print(relative_path)\n", + " \n", + " # intrinsics = intrinsic_for_all\n", + " pose = (np.array(frame['transform_matrix'])@blender2opencv)\n", + " # print(len(pose))\n", + " all_poses.append(pose)\n", + " \n", + " # cameras[relative_path] = {'pose': pose, 'intrinsics': intrinsics, 'scene-name': os.path.basename(scene_folder_path),\\\n", + " # 'pc_csv':pc_relative_path}\\\n", + "PREDEFINED_POSES = torch.tensor(np.stack(all_poses))\n", + "PREDEFINED_POSES\n", + "\n", + " " + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "B073WRF565\n", + "B07QHJZDH4\n", + "B07B4ZPGY3\n", + "B07MBFDRY5\n", + "B07DBHGJ15\n", + "B07B8RYY41\n", + "B07DBB1NJ1\n", + "B07TFFF6V4\n", + "B07B4YY8H3\n", + "B00WRDS8H0\n", + "B07MBFDWY5\n", + "B07DBCMTD3\n", + "B072LXZ653\n", + "B07TP5MKP3\n", + "B07DBHCKK5\n", + "B07GDZ8HN6\n", + "B07QC76NM2\n", + "B07QJM4QK6\n", + "B07W6NTK71\n", + "B07B78G443\n", + "B07QHMYBJ5\n", + "B072JC6WL3\n", + "B0825DWB31\n", + "B07B4Z6JZ4\n", + "B075Z96KX3\n", + "B075X2G512\n", + "B07QGHTCY4\n", + "B07BWJCVK3\n", + "B07B4GVRM3\n", + "B07CF7B5S2\n", + "B00IIFW2L4\n", + "B0853P1X94\n", + "B0719H3LG2\n", + "B075X2XZF4\n", + "B07B4VSTG2\n", + "B07ML7J584\n", + "B07MBF93P1\n", + "B07MF1RKM3\n", + "B07WBYQQ84\n", + "B07DBHB6D4\n", + "B071S5RLR6\n", + "B07MF1S1Z2\n", + "B07TF9MY62\n", + "B07MF1RQV5\n", + "B07K8V2SX3\n", + "B075X3S2Z6\n", + "B075X3S2Z5\n", + "B07DBJX741\n", + "B07JKTRKT3\n", + "B075X1TB14\n", + "B07MK6LHG2\n", + "B07HKGHTC1\n", + "B07HK85S41\n", + "B07QGG5263\n", + "B07F3XM8M3\n", + "B07QBQD983\n", + "B0716DKHS1\n", + "B07MBFDJQ3\n", + "B073772DH1\n", + "B07K7NQR23\n", + "B075X3S2Z1\n", + "B07DBCN4C2\n", + "B07B4DB6Z4\n", + "B07HKCJPJ3\n", + "B07BL1NYC4\n", + "B07HKFRV55\n", + "B075X3S2X5\n", + "B075HR1XB1\n", + "B075X2NLK5\n", + "B073NZT573\n", + "B075X2LLX3\n", + "B07MF1SN52\n", + "B07P5LLX43\n", + "B0742DNY41\n", + "B073G947W3\n", + "B07GFRKNR1\n", + "B07ML7M196\n", + "B07B4M6933\n", + "B0719FLQP3\n", + "B07DTLWMN5\n", + "B07B4Z9BS4\n", + "B07DT4GYP3\n", + "B07QHJZ794\n", + "B07B813LW1\n", + "B07B4XK5V3\n", + "B07DT153M1\n", + "B07GG1Z4J3\n", + "B07MBFDHP1\n", + "B07374K536\n", + "B07MBFG1D2\n", + "B07QJXMX23\n", + "B07HK3HJ15\n", + "B07HKGY4F2\n", + "B073P19B53\n", + "B07K4ZCQC3\n", + "B07HK67M71\n", + "B07MHMSDJ5\n", + "B07B4YHFR3\n", + "B07MBFF196\n", + "B07B4ZC524\n", + "B073P2N985\n", + "B07B4YHYR2\n", + "B07MHMJD43\n", + "B07B8P1J64\n", + "B07Q5SBTK5\n", + "B07JCDQWM6\n", + "B07MF1RNW1\n", + "B075HRFFQ3\n", + "B084RZVHD2\n" + ] + } + ], + "source": [ + "### filter those not completed data generation\n", + "\n", + "source = '/home/xuyi/Data/renderer/output_abo'\n", + "# completed_data_f = open('completed.txt', 'w')\n", + "train_data = 'abo_512_train.txt'\n", + "val_data = 'abo_512_val.txt'\n", + "train_f = open(train_data, 'w')\n", + "val_f = open(val_data, 'w')\n", + "all_data=os.listdir(source)\n", + "# val_i = all_data//\n", + "val_i=100\n", + "for i, folder in enumerate(all_data):\n", + " try:\n", + " sample_len=len(os.listdir(os.path.join(source,folder,'sample')))\n", + " render_len=len(os.listdir(os.path.join(source,folder,'render')))\n", + " if render_len <302 or sample_len < 2:\n", + " print(folder) \n", + " else:\n", + " # completed_data_f.writelines([f'{folder}\\n'])\n", + " if i%val_i == 0:\n", + " val_f.writelines([f'{folder}\\n'])\n", + " else:\n", + " train_f.writelines([f'{folder}\\n'])\n", + " pass\n", + " except:\n", + " pass\n", + "\n", + "# completed_data_f.close()\n", + "train_f.close()\n", + "val_f.close()" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "1351\n", + "14\n" + ] + }, + { + "data": { + "text/plain": [ + "1365" + ] + }, + "execution_count": 1, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# with open('completed.txt') as f:\n", + "# scans = [line.rstrip() for line in f.readlines()]\n", + "# print(len(scans))\n", + "\n", + "all_data = []\n", + "resolution=512\n", + "for split in ['train', 'val']:\n", + " with open(os.path.join('/home/xuyi/Data/renderer/output_abo', 'meta', f'abo_{resolution}_{split}.txt')) as f:\n", + " scans = [line.rstrip() for line in f.readlines()]\n", + " print(len(scans))\n", + " all_data += scans\n", + "len(all_data)\n" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "swapping_prob = 0.3\n", + "gen_c = torch.rand([5,10])\n", + "pc_dim=8\n", + "c_swapped = torch.roll(gen_c.clone(), 1, 0)\n", + "c_gen_conditioning = torch.where(torch.rand([], device=gen_c.device) < swapping_prob, c_swapped, gen_c)" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "torch.Size([5, 8])" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "gen_c = torch.rand([5,10])\n", + "pc_dim=8\n", + "gen_c[:,-8:].shape" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "(-0.4, 0.4)" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "data = pd.read_csv(csv_f)\n", + "data = data[['x','y','z','r','g','b','a', 'metallic','roughness']].values.astype(np.float32)\n", + "data.shape #numpy array\n", + "data[:,:3].min(), data[:,:3].max()" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + ">" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "import os\n", + "import io\n", + "import csv\n", + "import zipfile\n", + "from typing import Callable, Optional, Tuple, Union\n", + "pc_array = data\n", + "archive_fname = '00000/img00000000.png'\n", + "archive_root_dir = ''\n", + "dest = 'read_data.zip'\n", + "zf = zipfile.ZipFile(file=dest, mode='w', compression=zipfile.ZIP_STORED)\n", + "def zip_write_bytes(fname: str, data: Union[bytes, str]):\n", + " # st()\n", + " zf.writestr(fname, data)\n", + "def zip_write_csv(fname: str, data: Union[bytes, str]):\n", + " # st()\n", + " zf.writestr(fname, data)\n", + "\n", + "archive_fname_pc = archive_fname.replace('img', 'pc').replace('png', 'csv')\n", + "pc_fname = os.path.join(archive_root_dir, archive_fname_pc)\n", + "# os.makedirs(os.path.dirname(pc_fname), exist_ok=True)\n", + "# np.savetxt(pc_fname, pc_array, delimiter=\",\")\n", + "# pc_bits = io.BytesIO()\n", + "string_buffer = io.StringIO()\n", + "writer = csv.writer(string_buffer)\n", + "for row in pc_array:\n", + " # print(row)\n", + " writer.writerow(row)\n", + "zip_write_csv(os.path.join(archive_root_dir, archive_fname_pc), string_buffer.getvalue())\n", + "zf.close\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "'B07B4MHTG1', 'B07B4MF6P2', 'B07RTZ54B1', 'B075X4F3Z2', 'B07JL5QBC2', 'B075YPKYM1', 'B073NZGLT1', 'B07DYK2Y61'" + ] + }, + { + "cell_type": "code", + "execution_count": 34, + "metadata": {}, + "outputs": [], + "source": [ + "_path = '/home/xuyi/Repo/eg3d/dataset_preprocessing/abo/read_data.zip'\n", + "zf_reopen = zipfile.ZipFile(_path)\n", + "# os.path.isdir(_path)\n", + "_all_fnames = set(zf_reopen.namelist())\n", + "_all_fnames\n", + "fname = [i for i in _all_fnames][0]\n", + "fname\n", + "with zf_reopen.open(fname, 'r') as f:\n", + " # if pyspng is not None and self._file_ext(fname) == '.png':\n", + " # image = pyspng.load(f.read())\n", + " # else:\n", + " # image = np.array(PIL.Image.open(f))\n", + " df = pd.read_csv(f, header=None)\n", + " pc_array = df.values.astype(np.float32)\n" + ] + }, + { + "cell_type": "code", + "execution_count": 51, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "True" + ] + }, + "execution_count": 51, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "(df.values.astype(np.float32)).shape\n", + "\n", + "np.empty([3,3]).shape == (3,3)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "import os\n", + "def list_recursive(folderpath):\n", + " return [os.path.join(folderpath, filename) for filename in os.listdir(folderpath)]\n", + "dataset_path = '/home/xuyi/Data/renderer/output_abo'" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "B01IJ5A2UA/sample/pc.csv\n", + "B01N6AQX0A/sample/pc.csv\n", + "B01D3C7Z4A/sample/pc.csv\n" + ] + } + ], + "source": [ + "cameras = {}\n", + "blender2opencv = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])\n", + "\n", + "for scene_folder_path in list_recursive(dataset_path):\n", + " if not os.path.isdir(scene_folder_path): continue\n", + "\n", + " # st() # the sibling folder with rgb should be mesh>: no, only intrinsics and pose\n", + " \n", + " pointcloud_csv = os.path.join(scene_folder_path,'sample', f\"pc.csv\")\n", + " assert os.path.isfile(pointcloud_csv)\n", + " pc_relative_path = os.path.relpath(pointcloud_csv, dataset_path)\n", + " print(pc_relative_path)\n", + " \n", + " with open(os.path.join(scene_folder_path,'render', f\"transforms.json\"), 'r') as f:\n", + " meta = json.load(f)\n", + " # print(meta.keys()) ['camera_angle_x', 'frames']\n", + " # print(meta ['frames'][0]['file_path']) \n", + "\n", + " w, h = 512, 512\n", + " focal = .5 * w / np.tan(0.5 * meta['camera_angle_x'])\n", + " intrinsic_for_all = np.array([[focal, 0, w / 2], [0, focal, h / 2], [0, 0, 1]])\n", + "\n", + " # continue\n", + " # for rgb_path in list_recursive(os.path.join(scene_folder_path, 'render')):\n", + " for frame in meta ['frames']:\n", + " rgb_path = frame['file_path']\n", + " relative_path = os.path.relpath(rgb_path, dataset_path)\n", + " # print(relative_path)\n", + " \n", + " intrinsics = intrinsic_for_all.tolist()\n", + " pose = (np.array(frame['transform_matrix'])@blender2opencv).tolist()\n", + " # print(len(pose))\n", + " \n", + " \n", + " cameras[relative_path] = {'pose': pose, 'intrinsics': intrinsics, 'scene-name': os.path.basename(scene_folder_path),\\\n", + " 'pc_csv':pc_relative_path}\n", + " # if DEBUG:\n", + " # break" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "numpy.ndarray" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "data = pd.read_csv(csv)\n", + "vertices_pos = data[['x','y','z','r','g,']].values.astype(np.float32)\n", + "type(vertices_pos)\n" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [], + "source": [ + "dataset = {'labels':[]}\n", + "source = dataset_path\n", + "max_images = len(cameras)\n", + "for i, filename in enumerate(cameras):\n", + " if (max_images is not None and i >= max_images): break\n", + "\n", + " pose = np.array(cameras[filename]['pose'])\n", + " intrinsics = np.array(cameras[filename]['intrinsics'])\n", + " label = np.concatenate([pose.reshape(-1), intrinsics.reshape(-1)]).tolist()\n", + " \n", + " image_path = os.path.join(source, filename)\n", + " dataset[\"labels\"].append([filename, label])" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "[['B01IJ5A2UA/render/r_0',\n", + " [-1.0,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 0.9283531308174133,\n", + " -0.3716994822025299,\n", + " 0.6318891644477844,\n", + " 0.0,\n", + " -0.3716995418071747,\n", + " -0.9283530712127686,\n", + " 1.578200340270996,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_1',\n", + " [-0.9685832262039185,\n", + " -0.2308720350265503,\n", + " 0.09243789315223694,\n", + " -0.1571444272994995,\n", + " -0.24868987500667572,\n", + " 0.8991873264312744,\n", + " -0.3600218594074249,\n", + " 0.612037181854248,\n", + " 0.0,\n", + " -0.37169939279556274,\n", + " -0.9283530712127686,\n", + " 1.578200340270996,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_2',\n", + " [-0.8763066530227661,\n", + " -0.4467415511608124,\n", + " 0.18030120432376862,\n", + " -0.3065120577812195,\n", + " -0.4817536473274231,\n", + " 0.8126199245452881,\n", + " -0.32796669006347656,\n", + " 0.5575433969497681,\n", + " 0.0,\n", + " -0.374260276556015,\n", + " -0.9273236989974976,\n", + " 1.5764503479003906,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_3',\n", + " [-0.728968620300293,\n", + " -0.63266521692276,\n", + " 0.2614181935787201,\n", + " -0.44441089034080505,\n", + " -0.6845471262931824,\n", + " 0.6737200021743774,\n", + " -0.2783820927143097,\n", + " 0.47324952483177185,\n", + " 0.0,\n", + " -0.3818848431110382,\n", + " -0.9242098927497864,\n", + " 1.5711567401885986,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_4',\n", + " [-0.5358268022537231,\n", + " -0.7758848667144775,\n", + " 0.33300474286079407,\n", + " -0.5661081075668335,\n", + " -0.844327986240387,\n", + " 0.49239152669906616,\n", + " -0.21133123338222504,\n", + " 0.3592631220817566,\n", + " -1.4901162970204496e-08,\n", + " -0.39440208673477173,\n", + " -0.9189379215240479,\n", + " 1.562194585800171,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_5',\n", + " [-0.30901700258255005,\n", + " -0.8667895197868347,\n", + " 0.39138779044151306,\n", + " -0.6653592586517334,\n", + " -0.9510565400123596,\n", + " 0.2816369831562042,\n", + " -0.12716959416866302,\n", + " 0.21618832647800446,\n", + " -7.450581485102248e-09,\n", + " -0.4115295112133026,\n", + " -0.9113963842391968,\n", + " 1.549373984336853,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_6',\n", + " [-0.06279049813747406,\n", + " -0.8996737003326416,\n", + " 0.432023823261261,\n", + " -0.7344405651092529,\n", + " -0.9980267286300659,\n", + " 0.056602656841278076,\n", + " -0.0271806288510561,\n", + " 0.04620707035064697,\n", + " 1.862645149230957e-09,\n", + " -0.43287795782089233,\n", + " -0.9014524817466736,\n", + " 1.5324692726135254,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_7',\n", + " [0.18738128244876862,\n", + " -0.8732262849807739,\n", + " 0.44984909892082214,\n", + " -0.764743447303772,\n", + " -0.9822872281074524,\n", + " -0.1665768027305603,\n", + " 0.08581329882144928,\n", + " -0.14588260650634766,\n", + " -7.450580152834618e-09,\n", + " -0.45796069502830505,\n", + " -0.8889724016189575,\n", + " 1.511252999305725,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_8',\n", + " [0.42577940225601196,\n", + " -0.7906787991523743,\n", + " 0.43993061780929565,\n", + " -0.7478820085525513,\n", + " -0.9048269987106323,\n", + " -0.37206530570983887,\n", + " 0.20701570808887482,\n", + " -0.35192668437957764,\n", + " 0.0,\n", + " -0.48620399832725525,\n", + " -0.8738452792167664,\n", + " 1.4855369329452515,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_9',\n", + " [0.6374241709709167,\n", + " -0.6595649719238281,\n", + " 0.3983269929885864,\n", + " -0.6771558523178101,\n", + " -0.7705129981040955,\n", + " -0.5456399321556091,\n", + " 0.3295249342918396,\n", + " -0.560192346572876,\n", + " 1.4901160305669237e-08,\n", + " -0.5169631838798523,\n", + " -0.8560075759887695,\n", + " 1.4552128314971924,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_10',\n", + " [0.8090170621871948,\n", + " -0.49107447266578674,\n", + " 0.32301270961761475,\n", + " -0.549121618270874,\n", + " -0.5877850651741028,\n", + " -0.6759063601493835,\n", + " 0.4445890784263611,\n", + " -0.7558014988899231,\n", + " 0.0,\n", + " -0.549542248249054,\n", + " -0.8354659080505371,\n", + " 1.4202921390533447,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_11',\n", + " [0.9297766089439392,\n", + " -0.29903343319892883,\n", + " 0.21469616889953613,\n", + " -0.3649834990501404,\n", + " -0.36812418699264526,\n", + " -0.7552731037139893,\n", + " 0.5422612428665161,\n", + " -0.9218441247940063,\n", + " 0.0,\n", + " -0.5832167267799377,\n", + " -0.8123165965080261,\n", + " 1.3809382915496826,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_12',\n", + " [0.9921147227287292,\n", + " -0.0986066609621048,\n", + " 0.07736290246248245,\n", + " -0.1315169334411621,\n", + " -0.12533271312713623,\n", + " -0.7805553078651428,\n", + " 0.6123929619789124,\n", + " -1.0410679578781128,\n", + " 0.0,\n", + " -0.6172602772712708,\n", + " -0.7867591381072998,\n", + " 1.337490439414978,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_13',\n", + " [0.9921146631240845,\n", + " 0.09514132142066956,\n", + " -0.08158866316080093,\n", + " 0.1387007087469101,\n", + " 0.12533387541770935,\n", + " -0.7531171441078186,\n", + " 0.6458374261856079,\n", + " -1.0979235172271729,\n", + " 7.450580596923828e-09,\n", + " -0.650970458984375,\n", + " -0.759103000164032,\n", + " 1.2904748916625977,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_14',\n", + " [0.9297762513160706,\n", + " 0.2686455249786377,\n", + " -0.25168582797050476,\n", + " 0.4278659224510193,\n", + " 0.36812520027160645,\n", + " -0.6785197854042053,\n", + " 0.6356844902038574,\n", + " -1.0806636810302734,\n", + " -2.9802322387695312e-08,\n", + " -0.6836962699890137,\n", + " -0.7297666668891907,\n", + " 1.2406034469604492,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_15',\n", + " [0.8090165257453918,\n", + " 0.4110202491283417,\n", + " -0.42018404603004456,\n", + " 0.7143128514289856,\n", + " 0.5877858996391296,\n", + " -0.5657198429107666,\n", + " 0.5783327221870422,\n", + " -0.9831656217575073,\n", + " 0.0,\n", + " -0.7148589491844177,\n", + " -0.6992686986923218,\n", + " 1.1887567043304443,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_16',\n", + " [0.637423574924469,\n", + " 0.5148655772209167,\n", + " -0.5732404589653015,\n", + " 0.9745088815689087,\n", + " 0.7705137133598328,\n", + " -0.425933301448822,\n", + " 0.47422516345977783,\n", + " -0.8061828017234802,\n", + " 1.4901164746561335e-08,\n", + " -0.7439718842506409,\n", + " -0.6682108640670776,\n", + " 1.1359585523605347,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_17',\n", + " [0.42577868700027466,\n", + " 0.5766082406044006,\n", + " -0.6973059177398682,\n", + " 1.185420036315918,\n", + " 0.9048274159431458,\n", + " -0.27133074402809143,\n", + " 0.3281266391277313,\n", + " -0.5578152537345886,\n", + " -1.4901161193847656e-08,\n", + " -0.7706508040428162,\n", + " -0.6372576951980591,\n", + " 1.0833380222320557,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_18',\n", + " [0.18738046288490295,\n", + " 0.5963560342788696,\n", + " -0.780543327331543,\n", + " 1.3269237279891968,\n", + " 0.9822873473167419,\n", + " -0.11376047879457474,\n", + " 0.14889590442180634,\n", + " -0.25312304496765137,\n", + " -7.450580152834618e-09,\n", + " -0.7946180105209351,\n", + " -0.607109546661377,\n", + " 1.0320862531661987,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_19',\n", + " [-0.06279150396585464,\n", + " 0.5773343443870544,\n", + " -0.814089834690094,\n", + " 1.3839528560638428,\n", + " 0.9980267286300659,\n", + " 0.036323368549346924,\n", + " -0.051218997687101364,\n", + " 0.08707230538129807,\n", + " -1.862645149230957e-09,\n", + " -0.8156995177268982,\n", + " -0.5784757733345032,\n", + " 0.9834089279174805,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_20',\n", + " [-0.3090180456638336,\n", + " 0.5250285863876343,\n", + " -0.7930024266242981,\n", + " 1.3481042385101318,\n", + " 0.951056182384491,\n", + " 0.1705927550792694,\n", + " -0.2576630711555481,\n", + " 0.4380272328853607,\n", + " 0.0,\n", + " -0.833812415599823,\n", + " -0.5520479083061218,\n", + " 0.9384814500808716,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_21',\n", + " [-0.5358278155326843,\n", + " 0.4462044835090637,\n", + " -0.7167914509773254,\n", + " 1.2185455560684204,\n", + " 0.8443272709846497,\n", + " 0.2831707298755646,\n", + " -0.4548909068107605,\n", + " 0.7733145952224731,\n", + " 0.0,\n", + " -0.8489498496055603,\n", + " -0.5284733772277832,\n", + " 0.8984048366546631,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_22',\n", + " [-0.7289694547653198,\n", + " 0.3479785919189453,\n", + " -0.5895034074783325,\n", + " 1.0021559000015259,\n", + " 0.6845460534095764,\n", + " 0.37056055665016174,\n", + " -0.6277590394020081,\n", + " 1.067190408706665,\n", + " 1.4901161193847656e-08,\n", + " -0.8611595034599304,\n", + " -0.508334755897522,\n", + " 0.8641691207885742,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_23',\n", + " [-0.8763073086738586,\n", + " 0.23708508908748627,\n", + " -0.41937580704689026,\n", + " 0.7129388451576233,\n", + " 0.4817523956298828,\n", + " 0.4312576353549957,\n", + " -0.7628443241119385,\n", + " 1.2968353033065796,\n", + " 0.0,\n", + " -0.8705213665962219,\n", + " -0.4921306073665619,\n", + " 0.8366219997406006,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_24',\n", + " [-0.9685835838317871,\n", + " 0.11943555623292923,\n", + " -0.21813085675239563,\n", + " 0.3708224892616272,\n", + " 0.24868841469287872,\n", + " 0.46517375111579895,\n", + " -0.8495690226554871,\n", + " 1.4442673921585083,\n", + " 0.0,\n", + " -0.8771252036094666,\n", + " -0.48026183247566223,\n", + " 0.8164451718330383,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_25',\n", + " [-1.0,\n", + " -7.59368504077429e-07,\n", + " 1.4144011402095202e-06,\n", + " -2.4044818474067142e-06,\n", + " -1.6053569424911984e-06,\n", + " 0.47302156686782837,\n", + " -0.8810508847236633,\n", + " 1.4977864027023315,\n", + " 5.6843412084544437e-14,\n", + " -0.8810508251190186,\n", + " -0.47302162647247314,\n", + " 0.8041367530822754,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_26',\n", + " [-0.9685827493667603,\n", + " -0.11703130602836609,\n", + " 0.21943369507789612,\n", + " -0.37303730845451355,\n", + " -0.24869152903556824,\n", + " 0.45580363273620605,\n", + " -0.8546318411827087,\n", + " 1.4528741836547852,\n", + " 7.450581485102248e-09,\n", + " -0.882352888584137,\n", + " -0.47058820724487305,\n", + " 0.800000011920929,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_27',\n", + " [-0.8763058185577393,\n", + " -0.22788068652153015,\n", + " 0.42445090413093567,\n", + " -0.7215664982795715,\n", + " -0.4817553162574768,\n", + " 0.41451162099838257,\n", + " -0.7720699906349182,\n", + " 1.3125189542770386,\n", + " 1.4901161193847656e-08,\n", + " -0.8810508251190186,\n", + " -0.47302162647247314,\n", + " 0.8041367530822754,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_28',\n", + " [-0.7289673089981079,\n", + " -0.32876256108283997,\n", + " 0.6004347801208496,\n", + " -1.0207390785217285,\n", + " -0.684548556804657,\n", + " 0.3500951826572418,\n", + " -0.6393955945968628,\n", + " 1.0869724750518799,\n", + " 1.4901159417490817e-08,\n", + " -0.8771252036094666,\n", + " -0.480261892080307,\n", + " 0.8164451718330383,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_29',\n", + " [-0.5358250737190247,\n", + " -0.4155201017856598,\n", + " 0.7350064516067505,\n", + " -1.2495110034942627,\n", + " -0.8443290591239929,\n", + " 0.2636958956718445,\n", + " -0.4664471745491028,\n", + " 0.7929602265357971,\n", + " 0.0,\n", + " -0.8705213665962219,\n", + " -0.4921305477619171,\n", + " 0.8366219997406006,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_30',\n", + " [-0.3090149462223053,\n", + " -0.4834554195404053,\n", + " 0.819011926651001,\n", + " -1.3923203945159912,\n", + " -0.9510571956634521,\n", + " 0.15708304941654205,\n", + " -0.2661111652851105,\n", + " 0.45238903164863586,\n", + " 0.0,\n", + " -0.8611595034599304,\n", + " -0.508334755897522,\n", + " 0.8641691207885742,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_31',\n", + " [-0.06278830021619797,\n", + " -0.5274306535720825,\n", + " 0.8472747206687927,\n", + " -1.4403671026229858,\n", + " -0.998026967048645,\n", + " 0.03318193927407265,\n", + " -0.05330410972237587,\n", + " 0.09061699360609055,\n", + " 3.725290742551124e-09,\n", + " -0.8489498496055603,\n", + " -0.5284733772277832,\n", + " 0.8984048366546631,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_32',\n", + " [0.18738357722759247,\n", + " -0.5422694087028503,\n", + " 0.819042980670929,\n", + " -1.392372965812683,\n", + " -0.982286810874939,\n", + " -0.10344472527503967,\n", + " 0.15624277293682098,\n", + " -0.2656126916408539,\n", + " 0.0,\n", + " -0.833812415599823,\n", + " -0.5520479083061218,\n", + " 0.9384814500808716,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_33',\n", + " [0.4257814884185791,\n", + " -0.5234200358390808,\n", + " 0.7380661368370056,\n", + " -1.254712462425232,\n", + " -0.9048261046409607,\n", + " -0.24630430340766907,\n", + " 0.34730973839759827,\n", + " -0.5904265642166138,\n", + " -1.4901161193847656e-08,\n", + " -0.8156995177268982,\n", + " -0.578475832939148,\n", + " 0.9834089279174805,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_34',\n", + " [0.6374260187149048,\n", + " -0.46778497099876404,\n", + " 0.612262487411499,\n", + " -1.0408462285995483,\n", + " -0.7705116271972656,\n", + " -0.3869874179363251,\n", + " 0.5065102577209473,\n", + " -0.8610674142837524,\n", + " 0.0,\n", + " -0.7946180105209351,\n", + " -0.6071096062660217,\n", + " 1.0320862531661987,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_35',\n", + " [0.809018611907959,\n", + " -0.374569296836853,\n", + " 0.4529755413532257,\n", + " -0.7700583934783936,\n", + " -0.58778315782547,\n", + " -0.5155532956123352,\n", + " 0.6234707832336426,\n", + " -1.0599002838134766,\n", + " 2.9802322387695312e-08,\n", + " -0.7706507444381714,\n", + " -0.6372576951980591,\n", + " 1.0833380222320557,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_36',\n", + " [0.9297774434089661,\n", + " -0.24598316848278046,\n", + " 0.2738724648952484,\n", + " -0.4655831456184387,\n", + " -0.36812201142311096,\n", + " -0.62128746509552,\n", + " 0.6917283535003662,\n", + " -1.1759381294250488,\n", + " 0.0,\n", + " -0.7439718246459961,\n", + " -0.6682109236717224,\n", + " 1.1359585523605347,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_37',\n", + " [0.9921151399612427,\n", + " -0.08763962984085083,\n", + " 0.08959357440471649,\n", + " -0.1523090898990631,\n", + " -0.12533043324947357,\n", + " -0.6937549114227295,\n", + " 0.7092223763465881,\n", + " -1.205678105354309,\n", + " -7.450581485102248e-09,\n", + " -0.7148590087890625,\n", + " -0.6992685794830322,\n", + " 1.1887567043304443,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_38',\n", + " [0.9921143651008606,\n", + " 0.09146613627672195,\n", + " -0.08569183945655823,\n", + " 0.1456761360168457,\n", + " 0.12533612549304962,\n", + " -0.7240119576454163,\n", + " 0.6783047914505005,\n", + " -1.1531182527542114,\n", + " 0.0,\n", + " -0.6836962103843689,\n", + " -0.7297666668891907,\n", + " 1.2406034469604492,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_39',\n", + " [0.9297753572463989,\n", + " 0.279446542263031,\n", + " -0.23964010179042816,\n", + " 0.4073881506919861,\n", + " 0.36812737584114075,\n", + " -0.705795168876648,\n", + " 0.6052564382553101,\n", + " -1.0289359092712402,\n", + " 0.0,\n", + " -0.6509706377983093,\n", + " -0.7591028809547424,\n", + " 1.2904748916625977,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_40',\n", + " [0.809015154838562,\n", + " 0.4624474346637726,\n", + " -0.36281803250312805,\n", + " 0.6167906522750854,\n", + " 0.5877878665924072,\n", + " -0.6365000009536743,\n", + " 0.49937283992767334,\n", + " -0.8489338159561157,\n", + " 1.4901161193847656e-08,\n", + " -0.6172602772712708,\n", + " -0.7867591381072998,\n", + " 1.337490439414978,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_41',\n", + " [0.6374214887619019,\n", + " 0.6259022951126099,\n", + " -0.4493774175643921,\n", + " 0.7639416456222534,\n", + " 0.7705153226852417,\n", + " -0.5177879929542542,\n", + " 0.37175488471984863,\n", + " -0.6319833397865295,\n", + " 0.0,\n", + " -0.5832167267799377,\n", + " -0.8123165965080261,\n", + " 1.3809382915496826,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_42',\n", + " [0.4257762134075165,\n", + " 0.7559534311294556,\n", + " -0.49724146723747253,\n", + " 0.8453105688095093,\n", + " 0.9048284888267517,\n", + " -0.3557215631008148,\n", + " 0.23398201167583466,\n", + " -0.3977694511413574,\n", + " 0.0,\n", + " -0.5495421886444092,\n", + " -0.8354659080505371,\n", + " 1.4202921390533447,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_43',\n", + " [0.1873779147863388,\n", + " 0.8408457636833191,\n", + " -0.5078068375587463,\n", + " 0.8632715940475464,\n", + " 0.9822879433631897,\n", + " -0.16039688885211945,\n", + " 0.09686751663684845,\n", + " -0.164674773812294,\n", + " 7.450580596923828e-09,\n", + " -0.5169633626937866,\n", + " -0.8560075759887695,\n", + " 1.4552128314971924,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_44',\n", + " [-0.06279406696557999,\n", + " 0.87212073802948,\n", + " -0.4852445125579834,\n", + " 0.8249157071113586,\n", + " 0.9980265498161316,\n", + " 0.0548722967505455,\n", + " -0.03053073026239872,\n", + " 0.05190224200487137,\n", + " -1.862645149230957e-09,\n", + " -0.4862040579319,\n", + " -0.8738452196121216,\n", + " 1.4855369329452515,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_45',\n", + " [-0.309020459651947,\n", + " 0.8454619646072388,\n", + " -0.4355461001396179,\n", + " 0.7404283285140991,\n", + " 0.9510554075241089,\n", + " 0.27471065521240234,\n", + " -0.14151926338672638,\n", + " 0.24058274924755096,\n", + " 0.0,\n", + " -0.45796069502830505,\n", + " -0.8889724016189575,\n", + " 1.511252999305725,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_46',\n", + " [-0.5358299612998962,\n", + " 0.7611197233200073,\n", + " -0.365490198135376,\n", + " 0.6213333606719971,\n", + " 0.8443259000778198,\n", + " 0.4830252528190613,\n", + " -0.23194904625415802,\n", + " 0.39431339502334595,\n", + " 0.0,\n", + " -0.4328780472278595,\n", + " -0.9014524817466736,\n", + " 1.5324692726135254,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_47',\n", + " [-0.7289713025093079,\n", + " 0.6238912343978882,\n", + " -0.2817101776599884,\n", + " 0.4789073169231415,\n", + " 0.6845443248748779,\n", + " 0.6643818020820618,\n", + " -0.2999931871891022,\n", + " 0.5099884271621704,\n", + " 0.0,\n", + " -0.41152945160865784,\n", + " -0.9113963842391968,\n", + " 1.549373984336853,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_48',\n", + " [-0.8763085603713989,\n", + " 0.44269853830337524,\n", + " -0.19000335037708282,\n", + " 0.32300570607185364,\n", + " 0.4817502498626709,\n", + " 0.8052731156349182,\n", + " -0.34561800956726074,\n", + " 0.5875506401062012,\n", + " 0.0,\n", + " -0.3944022059440613,\n", + " -0.9189379215240479,\n", + " 1.562194585800171,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01IJ5A2UA/render/r_49',\n", + " [-0.9685842394828796,\n", + " 0.229838028550148,\n", + " -0.09496940672397614,\n", + " 0.16144798696041107,\n", + " 0.24868600070476532,\n", + " 0.895175039768219,\n", + " -0.3698876202106476,\n", + " 0.6288089156150818,\n", + " 0.0,\n", + " -0.3818848729133606,\n", + " -0.9242098927497864,\n", + " 1.5711567401885986,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01N6AQX0A/render/r_0',\n", + " [-1.0,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 0.6442176699638367,\n", + " -0.7648422122001648,\n", + " 1.3002318143844604,\n", + " 0.0,\n", + " -0.7648422718048096,\n", + " -0.6442176103591919,\n", + " 1.095170021057129,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01N6AQX0A/render/r_1',\n", + " [0.80901700258255,\n", + " -0.378661572933197,\n", + " 0.4495629370212555,\n", + " -0.7642569541931152,\n", + " -0.5877851843833923,\n", + " -0.5211830735206604,\n", + " 0.6187704205513,\n", + " -1.0519096851348877,\n", + " 1.4901159417490817e-08,\n", + " -0.7648422718048096,\n", + " -0.6442176699638367,\n", + " 1.095170021057129,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01N6AQX0A/render/r_2',\n", + " [-0.3090171217918396,\n", + " 0.42064744234085083,\n", + " -0.8529736995697021,\n", + " 1.4500552415847778,\n", + " 0.9510564804077148,\n", + " 0.13667671382427216,\n", + " -0.2771480977535248,\n", + " 0.471151739358902,\n", + " 0.0,\n", + " -0.8968695998191833,\n", + " -0.44229501485824585,\n", + " 0.7519015073776245,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01N6AQX0A/render/r_3',\n", + " [-0.3090168237686157,\n", + " -0.06352514028549194,\n", + " 0.9489325881004333,\n", + " -1.6131855249404907,\n", + " -0.9510566592216492,\n", + " 0.020640553906559944,\n", + " -0.3083266615867615,\n", + " 0.5241553783416748,\n", + " 0.0,\n", + " -0.9977666735649109,\n", + " -0.06679428368806839,\n", + " 0.11355028301477432,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01N6AQX0A/render/r_4',\n", + " [0.8090168833732605,\n", + " 0.03926071152091026,\n", + " -0.5864728093147278,\n", + " 0.9970037937164307,\n", + " 0.5877854228019714,\n", + " -0.05403769761323929,\n", + " 0.8072100877761841,\n", + " -1.3722572326660156,\n", + " 0.0,\n", + " -0.9977666735649109,\n", + " -0.06679428368806839,\n", + " 0.11355028301477432,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_0',\n", + " [-1.0,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 0.8104843497276306,\n", + " -0.585760235786438,\n", + " 1.2075755596160889,\n", + " 0.0,\n", + " -0.5857601761817932,\n", + " -0.8104844689369202,\n", + " 1.6708564758300781,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_1',\n", + " [-0.9685831665992737,\n", + " -0.20155929028987885,\n", + " 0.14567264914512634,\n", + " -0.3003118336200714,\n", + " -0.24868988990783691,\n", + " 0.7850216031074524,\n", + " -0.5673574805259705,\n", + " 1.1696373224258423,\n", + " -7.4505797087454084e-09,\n", + " -0.5857602953910828,\n", + " -0.8104844689369202,\n", + " 1.6708564758300781,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_2',\n", + " [-0.8763067126274109,\n", + " -0.38967350125312805,\n", + " 0.2832685708999634,\n", + " -0.583973228931427,\n", + " -0.4817536473274231,\n", + " 0.7088135480880737,\n", + " -0.5152636766433716,\n", + " 1.0622434616088867,\n", + " 0.0,\n", + " -0.5879946351051331,\n", + " -0.8088646531105042,\n", + " 1.6675174236297607,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_3',\n", + " [-0.728968620300293,\n", + " -0.5503721833229065,\n", + " 0.4070567488670349,\n", + " -0.8391689658164978,\n", + " -0.6845471262931824,\n", + " 0.5860868096351624,\n", + " -0.4334714114665985,\n", + " 0.8936241865158081,\n", + " 1.4901161193847656e-08,\n", + " -0.5946363806724548,\n", + " -0.8039946556091309,\n", + " 1.6574773788452148,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_4',\n", + " [-0.5358268022537231,\n", + " -0.6719534397125244,\n", + " 0.5112417936325073,\n", + " -1.0539520978927612,\n", + " -0.844327986240387,\n", + " 0.42643457651138306,\n", + " -0.32444387674331665,\n", + " 0.6688582301139832,\n", + " 0.0,\n", + " -0.6055012941360474,\n", + " -0.7958441376686096,\n", + " 1.6406748294830322,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_5',\n", + " [-0.30901697278022766,\n", + " -0.7459839582443237,\n", + " 0.5899293422698975,\n", + " -1.2161704301834106,\n", + " -0.9510565996170044,\n", + " 0.24238485097885132,\n", + " -0.19167964160442352,\n", + " 0.39515769481658936,\n", + " 0.0,\n", + " -0.6202883124351501,\n", + " -0.7843738794326782,\n", + " 1.6170281171798706,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_6',\n", + " [-0.06279050558805466,\n", + " -0.7680305242538452,\n", + " 0.6373276710510254,\n", + " -1.3138847351074219,\n", + " -0.9980267882347107,\n", + " 0.04832036793231964,\n", + " -0.04009724408388138,\n", + " 0.08266258984804153,\n", + " 0.0,\n", + " -0.6385876536369324,\n", + " -0.7695489525794983,\n", + " 1.5864659547805786,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_7',\n", + " [0.1873812973499298,\n", + " -0.7380493879318237,\n", + " 0.6482062339782715,\n", + " -1.3363113403320312,\n", + " -0.9822872281074524,\n", + " -0.1407904475927353,\n", + " 0.12365194410085678,\n", + " -0.2549149990081787,\n", + " 0.0,\n", + " -0.6598947644233704,\n", + " -0.7513580918312073,\n", + " 1.548964262008667,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_8',\n", + " [0.42577940225601196,\n", + " -0.6603721976280212,\n", + " 0.618563175201416,\n", + " -1.275200605392456,\n", + " -0.9048269987106323,\n", + " -0.3107476532459259,\n", + " 0.29107382893562317,\n", + " -0.6000640392303467,\n", + " -1.4901161193847656e-08,\n", + " -0.6836258172988892,\n", + " -0.729832649230957,\n", + " 1.5045884847640991,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_9',\n", + " [0.637424111366272,\n", + " -0.5432634949684143,\n", + " 0.5464019775390625,\n", + " -1.126436471939087,\n", + " -0.770513117313385,\n", + " -0.4494268596172333,\n", + " 0.4520232379436493,\n", + " -0.9318697452545166,\n", + " -1.4901161193847656e-08,\n", + " -0.7091403007507324,\n", + " -0.7050672769546509,\n", + " 1.453533411026001,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_10',\n", + " [0.8090171813964844,\n", + " -0.39806994795799255,\n", + " 0.432471364736557,\n", + " -0.8915625214576721,\n", + " -0.5877849459648132,\n", + " -0.547896683216095,\n", + " 0.5952461957931519,\n", + " -1.2271313667297363,\n", + " 0.0,\n", + " -0.7357645034790039,\n", + " -0.6772374510765076,\n", + " 1.3961607217788696,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_11',\n", + " [0.9297766089439392,\n", + " -0.2380335032939911,\n", + " 0.28081214427948,\n", + " -0.5789090394973755,\n", + " -0.3681241571903229,\n", + " -0.6012046933174133,\n", + " 0.709251344203949,\n", + " -1.4621590375900269,\n", + " 0.0,\n", + " -0.7628191113471985,\n", + " -0.6466119885444641,\n", + " 1.3330247402191162,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_12',\n", + " [0.992114782333374,\n", + " -0.07689925283193588,\n", + " 0.09896866977214813,\n", + " -0.2040291279554367,\n", + " -0.12533271312713623,\n", + " -0.6087228059768677,\n", + " 0.7834209203720093,\n", + " -1.6150635480880737,\n", + " 3.725290076417309e-09,\n", + " -0.7896474599838257,\n", + " -0.6135608553886414,\n", + " 1.264888048171997,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_13',\n", + " [0.9921146631240845,\n", + " 0.07251245528459549,\n", + " -0.10222776979207993,\n", + " 0.21074794232845306,\n", + " 0.12533386051654816,\n", + " -0.5739923119544983,\n", + " 0.8092120885848999,\n", + " -1.6682333946228027,\n", + " 0.0,\n", + " -0.8156436085700989,\n", + " -0.5785545110702515,\n", + " 1.192720651626587,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_14',\n", + " [0.9297761917114258,\n", + " 0.19958168268203735,\n", + " -0.30932721495628357,\n", + " 0.6376944184303284,\n", + " 0.36812520027160645,\n", + " -0.5040847659111023,\n", + " 0.7812697291374207,\n", + " -1.6106289625167847,\n", + " -1.4901160305669237e-08,\n", + " -0.8402772545814514,\n", + " -0.542156994342804,\n", + " 1.117685317993164,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_15',\n", + " [0.8090164661407471,\n", + " 0.29683858156204224,\n", + " -0.5073254704475403,\n", + " 1.0458781719207764,\n", + " 0.5877858996391296,\n", + " -0.4085625410079956,\n", + " 0.6982724070549011,\n", + " -1.4395253658294678,\n", + " 1.4901160305669237e-08,\n", + " -0.8631126284599304,\n", + " -0.5050114393234253,\n", + " 1.0411076545715332,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_16',\n", + " [0.6374235153198242,\n", + " 0.3604617714881897,\n", + " -0.6809983253479004,\n", + " 1.403913974761963,\n", + " 0.7705137133598328,\n", + " -0.29819950461387634,\n", + " 0.5633700489997864,\n", + " -1.161417007446289,\n", + " 0.0,\n", + " -0.8838236927986145,\n", + " -0.46782010793685913,\n", + " 0.9644358158111572,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_17',\n", + " [0.4257785975933075,\n", + " 0.39027056097984314,\n", + " -0.8163341879844666,\n", + " 1.6829159259796143,\n", + " 0.9048274159431458,\n", + " -0.183647021651268,\n", + " 0.3841369152069092,\n", + " -0.7919185161590576,\n", + " -1.4901161193847656e-08,\n", + " -0.9021987915039062,\n", + " -0.43132051825523376,\n", + " 0.8891900181770325,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_18',\n", + " [0.18738046288490295,\n", + " 0.3892407715320587,\n", + " -0.901875913143158,\n", + " 1.859264850616455,\n", + " 0.9822874069213867,\n", + " -0.07425129413604736,\n", + " 0.17204120755195618,\n", + " -0.35467204451560974,\n", + " -7.450580152834618e-09,\n", + " -0.9181385040283203,\n", + " -0.39625951647758484,\n", + " 0.8169099688529968,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_19',\n", + " [-0.06279151886701584,\n", + " 0.3626507818698883,\n", + " -0.9298073053359985,\n", + " 1.9168468713760376,\n", + " 0.9980267286300659,\n", + " 0.02281641587615013,\n", + " -0.05849944427609444,\n", + " 0.12059969455003738,\n", + " 0.0,\n", + " -0.9316458106040955,\n", + " -0.36336779594421387,\n", + " 0.7491019368171692,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_20',\n", + " [-0.309018075466156,\n", + " 0.3170211911201477,\n", + " -0.8966635465621948,\n", + " 1.848519206047058,\n", + " 0.9510562419891357,\n", + " 0.10300680994987488,\n", + " -0.2913447320461273,\n", + " 0.6006225347518921,\n", + " -7.450580596923828e-09,\n", + " -0.9428081512451172,\n", + " -0.3333359360694885,\n", + " 0.6871895790100098,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_21',\n", + " [-0.5358278751373291,\n", + " 0.25903451442718506,\n", + " -0.8036103844642639,\n", + " 1.65668523311615,\n", + " 0.8443273305892944,\n", + " 0.1643887609243393,\n", + " -0.5099880695343018,\n", + " 1.0513672828674316,\n", + " 0.0,\n", + " -0.9517759084701538,\n", + " -0.3067939877510071,\n", + " 0.6324719786643982,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_22',\n", + " [-0.7289694547653198,\n", + " 0.1946132928133011,\n", + " -0.6562995910644531,\n", + " 1.3529962301254272,\n", + " 0.6845461130142212,\n", + " 0.20724263787269592,\n", + " -0.6988899111747742,\n", + " 1.4407984018325806,\n", + " -1.4901158529312397e-08,\n", + " -0.958736777305603,\n", + " -0.28429532051086426,\n", + " 0.5860897898674011,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_23',\n", + " [-0.8763073682785034,\n", + " 0.12829279899597168,\n", + " -0.4643558859825134,\n", + " 0.9572941660881042,\n", + " 0.4817524254322052,\n", + " 0.23336449265480042,\n", + " -0.8446630239486694,\n", + " 1.7413173913955688,\n", + " -1.4901160305669237e-08,\n", + " -0.9638890027999878,\n", + " -0.2663043737411499,\n", + " 0.5490005016326904,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_24',\n", + " [-0.9685836434364319,\n", + " 0.06296500563621521,\n", + " -0.24058540165424347,\n", + " 0.49597954750061035,\n", + " 0.2486884593963623,\n", + " 0.24523404240608215,\n", + " -0.9370241165161133,\n", + " 1.9317249059677124,\n", + " 3.725291186640334e-09,\n", + " -0.9674169421195984,\n", + " -0.2531883418560028,\n", + " 0.5219611525535583,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_25',\n", + " [-1.0,\n", + " -3.9365298221127887e-07,\n", + " 1.5563446140731685e-06,\n", + " -3.208486532457755e-06,\n", + " -1.6053570561780361e-06,\n", + " 0.245212122797966,\n", + " -0.9694694876670837,\n", + " 1.998612403869629,\n", + " 0.0,\n", + " -0.9694693684577942,\n", + " -0.2452121526002884,\n", + " 0.5055177807807922,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_26',\n", + " [-0.9685827493667603,\n", + " -0.0603165477514267,\n", + " 0.24126625061035156,\n", + " -0.49738308787345886,\n", + " -0.24869155883789062,\n", + " 0.23491577804088593,\n", + " -0.9396633505821228,\n", + " 1.9371654987335205,\n", + " 0.0,\n", + " -0.9701424837112427,\n", + " -0.24253563582897186,\n", + " 0.5,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_27',\n", + " [-0.8763058185577393,\n", + " -0.11813223361968994,\n", + " 0.4670470356941223,\n", + " -0.9628421068191528,\n", + " -0.4817552864551544,\n", + " 0.2148808091878891,\n", + " -0.8495517373085022,\n", + " 1.7513957023620605,\n", + " 0.0,\n", + " -0.969469428062439,\n", + " -0.2452121526002884,\n", + " 0.5055177807807922,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_28',\n", + " [-0.7289673686027527,\n", + " -0.17331969738006592,\n", + " 0.6622439026832581,\n", + " -1.3652507066726685,\n", + " -0.6845484972000122,\n", + " 0.18456603586673737,\n", + " -0.7052154541015625,\n", + " 1.45383882522583,\n", + " 0.0,\n", + " -0.9674169421195984,\n", + " -0.2531883716583252,\n", + " 0.5219611525535583,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_29',\n", + " [-0.5358250737190247,\n", + " -0.22484850883483887,\n", + " 0.8138394951820374,\n", + " -1.6777729988098145,\n", + " -0.8443291187286377,\n", + " 0.14269255101680756,\n", + " -0.5164759159088135,\n", + " 1.0647423267364502,\n", + " 7.450580596923828e-09,\n", + " -0.9638890027999878,\n", + " -0.2663043737411499,\n", + " 0.5490005016326904,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_30',\n", + " [-0.3090149760246277,\n", + " -0.2703810930252075,\n", + " 0.9118135571479797,\n", + " -1.8797516822814941,\n", + " -0.9510571360588074,\n", + " 0.08785151690244675,\n", + " -0.29626405239105225,\n", + " 0.6107639670372009,\n", + " 0.0,\n", + " -0.958736777305603,\n", + " -0.28429532051086426,\n", + " 0.5860897898674011,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_31',\n", + " [-0.06278830766677856,\n", + " -0.3061886131763458,\n", + " 0.9498980045318604,\n", + " -1.9582648277282715,\n", + " -0.998026967048645,\n", + " 0.019263070076704025,\n", + " -0.05976039543747902,\n", + " 0.1231992095708847,\n", + " 1.862645371275562e-09,\n", + " -0.9517759084701538,\n", + " -0.3067939877510071,\n", + " 0.6324719786643982,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_32',\n", + " [0.18738357722759247,\n", + " -0.32743147015571594,\n", + " 0.9261080622673035,\n", + " -1.909220576286316,\n", + " -0.982286810874939,\n", + " -0.06246167793869972,\n", + " 0.1766667664051056,\n", + " -0.3642078638076782,\n", + " 3.7252898543727042e-09,\n", + " -0.9428081512451172,\n", + " -0.3333359360694885,\n", + " 0.6871895790100098,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_33',\n", + " [0.4257814884185791,\n", + " -0.32878467440605164,\n", + " 0.8429772853851318,\n", + " -1.7378422021865845,\n", + " -0.9048260450363159,\n", + " -0.15471531450748444,\n", + " 0.39667749404907227,\n", + " -0.8177716135978699,\n", + " 1.4901161193847656e-08,\n", + " -0.9316457509994507,\n", + " -0.36336779594421387,\n", + " 0.7491019368171692,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_34',\n", + " [0.63742595911026,\n", + " -0.3053225874900818,\n", + " 0.7074364423751831,\n", + " -1.458417534828186,\n", + " -0.7705116868019104,\n", + " -0.25258609652519226,\n", + " 0.5852453112602234,\n", + " -1.2065141201019287,\n", + " 0.0,\n", + " -0.9181385040283203,\n", + " -0.3962595462799072,\n", + " 0.8169099688529968,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_35',\n", + " [0.809018611907959,\n", + " -0.2535228729248047,\n", + " 0.5302972197532654,\n", + " -1.093235731124878,\n", + " -0.58778315782547,\n", + " -0.3489462435245514,\n", + " 0.7298955917358398,\n", + " -1.5047181844711304,\n", + " 0.0,\n", + " -0.9021987318992615,\n", + " -0.43132051825523376,\n", + " 0.8891900181770325,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_36',\n", + " [0.9297775626182556,\n", + " -0.17221486568450928,\n", + " 0.3253549635410309,\n", + " -0.6707364320755005,\n", + " -0.36812207102775574,\n", + " -0.4349685609340668,\n", + " 0.8217594027519226,\n", + " -1.6941003799438477,\n", + " -1.4901162970204496e-08,\n", + " -0.8838236331939697,\n", + " -0.46782010793685913,\n", + " 0.9644358158111572,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_37',\n", + " [0.9921151399612427,\n", + " -0.06329328566789627,\n", + " 0.10817427933216095,\n", + " -0.2230069786310196,\n", + " -0.12533041834831238,\n", + " -0.501029372215271,\n", + " 0.8563071489334106,\n", + " -1.7653223276138306,\n", + " 0.0,\n", + " -0.8631126880645752,\n", + " -0.5050114393234253,\n", + " 1.0411076545715332,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_38',\n", + " [0.9921143054962158,\n", + " 0.06795185804367065,\n", + " -0.10531709343194962,\n", + " 0.21711677312850952,\n", + " 0.12533612549304962,\n", + " -0.5378817319869995,\n", + " 0.8336510062217712,\n", + " -1.7186157703399658,\n", + " -3.725290076417309e-09,\n", + " -0.8402771949768066,\n", + " -0.542156994342804,\n", + " 1.117685317993164,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_39',\n", + " [0.9297753572463989,\n", + " 0.21298174560070038,\n", + " -0.3002608120441437,\n", + " 0.6190034747123718,\n", + " 0.36812740564346313,\n", + " -0.537925660610199,\n", + " 0.7583654522895813,\n", + " -1.5634104013442993,\n", + " 0.0,\n", + " -0.8156436085700989,\n", + " -0.5785545110702515,\n", + " 1.192720651626587,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_40',\n", + " [0.809015154838562,\n", + " 0.36064356565475464,\n", + " -0.46414515376091003,\n", + " 0.9568597078323364,\n", + " 0.5877878069877625,\n", + " -0.4963800013065338,\n", + " 0.6388368010520935,\n", + " -1.3169957399368286,\n", + " 1.4901159417490817e-08,\n", + " -0.7896475195884705,\n", + " -0.6135608553886414,\n", + " 1.264888048171997,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_41',\n", + " [0.6374215483665466,\n", + " 0.4982243776321411,\n", + " -0.587763786315918,\n", + " 1.2117060422897339,\n", + " 0.7705153226852417,\n", + " -0.4121643602848053,\n", + " 0.4862373173236847,\n", + " -1.0024038553237915,\n", + " 0.0,\n", + " -0.7628190517425537,\n", + " -0.6466119885444641,\n", + " 1.3330247402191162,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_42',\n", + " [0.4257762134075165,\n", + " 0.6127837300300598,\n", + " -0.6657407283782959,\n", + " 1.3724596500396729,\n", + " 0.9048284888267517,\n", + " -0.2883515954017639,\n", + " 0.31327107548713684,\n", + " -0.6458248496055603,\n", + " 1.4901160305669237e-08,\n", + " -0.7357644438743591,\n", + " -0.6772374510765076,\n", + " 1.3961607217788696,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_43',\n", + " [0.1873779296875,\n", + " 0.6925790309906006,\n", + " -0.6965800523757935,\n", + " 1.436036467552185,\n", + " 0.9822878837585449,\n", + " -0.13211405277252197,\n", + " 0.13287727534770966,\n", + " -0.2739335000514984,\n", + " 0.0,\n", + " -0.7091403603553772,\n", + " -0.7050672769546509,\n", + " 1.453533411026001,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_44',\n", + " [-0.06279407441616058,\n", + " 0.7283922433853149,\n", + " -0.6822766661643982,\n", + " 1.4065494537353516,\n", + " 0.9980266094207764,\n", + " 0.045829154551029205,\n", + " -0.04292764514684677,\n", + " 0.08849760890007019,\n", + " 0.0,\n", + " -0.6836256980895996,\n", + " -0.7298325896263123,\n", + " 1.5045884847640991,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_45',\n", + " [-0.309020459651947,\n", + " 0.7145830988883972,\n", + " -0.6275964379310608,\n", + " 1.2938231229782104,\n", + " 0.9510554075241089,\n", + " 0.2321849912405014,\n", + " -0.20392096042633057,\n", + " 0.42039382457733154,\n", + " 0.0,\n", + " -0.6598946452140808,\n", + " -0.7513580918312073,\n", + " 1.548964262008667,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_46',\n", + " [-0.5358299612998962,\n", + " 0.6497502326965332,\n", + " -0.5391762256622314,\n", + " 1.1115401983261108,\n", + " 0.8443259000778198,\n", + " 0.4123474359512329,\n", + " -0.3421744704246521,\n", + " 0.7054107189178467,\n", + " -1.4901159417490817e-08,\n", + " -0.6385877728462219,\n", + " -0.7695490717887878,\n", + " 1.5864659547805786,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_47',\n", + " [-0.7289713025093079,\n", + " 0.5369386672973633,\n", + " -0.42461487650871277,\n", + " 0.8753659725189209,\n", + " 0.6845442652702332,\n", + " 0.571786105632782,\n", + " -0.4521724283695221,\n", + " 0.9321773052215576,\n", + " 1.4901159417490817e-08,\n", + " -0.6202883124351501,\n", + " -0.7843738794326782,\n", + " 1.6170281171798706,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_48',\n", + " [-0.8763085603713989,\n", + " 0.38339805603027344,\n", + " -0.29170048236846924,\n", + " 0.6013559699058533,\n", + " 0.4817502498626709,\n", + " 0.6974049210548401,\n", + " -0.5306061506271362,\n", + " 1.0938726663589478,\n", + " -1.4901162970204496e-08,\n", + " -0.6055014729499817,\n", + " -0.7958441376686096,\n", + " 1.6406748294830322,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]],\n", + " ['B01D3C7Z4A/render/r_49',\n", + " [-0.9685841798782349,\n", + " 0.19994214177131653,\n", + " -0.1478777527809143,\n", + " 0.3048577904701233,\n", + " 0.24868595600128174,\n", + " 0.778736412525177,\n", + " -0.5759555697441101,\n", + " 1.187362790107727,\n", + " 0.0,\n", + " -0.5946366786956787,\n", + " -0.8039946556091309,\n", + " 1.6574773788452148,\n", + " 0.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0,\n", + " 711.1110599640117,\n", + " 0.0,\n", + " 256.0,\n", + " 0.0,\n", + " 711.1110599640117,\n", + " 256.0,\n", + " 0.0,\n", + " 0.0,\n", + " 1.0]]]" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3.9.13 ('eg3d')", + "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.13 | packaged by conda-forge | (main, May 27 2022, 16:56:21) \n[GCC 10.3.0]" + }, + "orig_nbformat": 4, + "vscode": { + "interpreter": { + "hash": "b417a2e9f2e73df9c6cd058b852a57ef19cc418a83dde0961c10fe9fe6466d1d" + } + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/dataset_preprocessing/abo/read_data.zip b/dataset_preprocessing/abo/read_data.zip new file mode 100644 index 00000000..9ddec388 Binary files /dev/null and b/dataset_preprocessing/abo/read_data.zip differ diff --git a/dataset_preprocessing/abo/run_me.py b/dataset_preprocessing/abo/run_me.py new file mode 100644 index 00000000..2bfbae4d --- /dev/null +++ b/dataset_preprocessing/abo/run_me.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +import os +import gdown +import shutil +import tempfile +import subprocess + + +if __name__ == '__main__': + with tempfile.TemporaryDirectory() as working_dir: + # working_dir = '/tmp/tmphal02_sj' # /cars_train.zip + working_dir = '/home/xuyi/Data' + # print(working_dir) + # download_name = 'cars_train.zip' + # url = 'https://drive.google.com/uc?id=1bThUNtIHx4xEQyffVBSf82ABDDh2HlFn' + # output_dataset_name = 'abo_128_completed.zip' + # output_dataset_name = 'abo_128_completed_white.zip' + # output_dataset_name = 'abo_512_completed_white.zip' + output_dataset_name = 'debug.zip' + + dir_path = os.path.dirname(os.path.realpath(__file__)) + # extracted_data_path = os.path.join(working_dir, os.path.splitext(download_name)[0]) + extracted_data_path = '/home/xuyi/Data/renderer/output_abo' + print("Downloading data...") + # zipped_dataset = os.path.join(working_dir, download_name) + # gdown.download(url, zipped_dataset, quiet=False) + + print("Unzipping downloaded data...") + # shutil.unpack_archive(zipped_dataset, working_dir) + + print("Converting camera parameters...") + cmd = f"python {os.path.join(dir_path, 'preprocess_abo_cameras.py')} --source={extracted_data_path}" + subprocess.run([cmd], shell=True) + + print("Creating dataset zip...") + cmd = f"python {os.path.join(dir_path, '../../eg3d', 'dataset_tool.py')}" + cmd += f" --source {extracted_data_path} --dest {output_dataset_name} --resolution 512x512 --read_pointcloud" + subprocess.run([cmd], shell=True) \ No newline at end of file diff --git a/dataset_preprocessing/shapenet_cars/preprocess_shapenet_cameras.py b/dataset_preprocessing/shapenet_cars/preprocess_shapenet_cameras.py index 80cbcb82..756d06cb 100644 --- a/dataset_preprocessing/shapenet_cars/preprocess_shapenet_cameras.py +++ b/dataset_preprocessing/shapenet_cars/preprocess_shapenet_cameras.py @@ -15,15 +15,19 @@ ############################################################# +# from distutils.debug import DEBUG import json import numpy as np import os from tqdm import tqdm import argparse +from ipdb import set_trace as st def list_recursive(folderpath): return [os.path.join(folderpath, filename) for filename in os.listdir(folderpath)] +DEBUG=False + if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument("--source", type=str) @@ -35,6 +39,9 @@ def list_recursive(folderpath): cameras = {} for scene_folder_path in list_recursive(dataset_path): if not os.path.isdir(scene_folder_path): continue + + # st() # the sibling folder with rgb should be mesh>: no, only intrinsics and pose + for rgb_path in list_recursive(os.path.join(scene_folder_path, 'rgb')): relative_path = os.path.relpath(rgb_path, dataset_path) @@ -53,7 +60,9 @@ def list_recursive(folderpath): cx = float(first_line[1]) cy = float(first_line[2]) + ### FIXME: need to adjust this "orig_img_size" if not using 512 image orig_img_size = 512 # cars_train has intrinsics corresponding to image size of 512 * 512 + intrinsics = np.array( [[focal / orig_img_size, 0.00000000e+00, cx / orig_img_size], [0.00000000e+00, focal / orig_img_size, cy / orig_img_size], @@ -61,6 +70,8 @@ def list_recursive(folderpath): ).tolist() cameras[relative_path] = {'pose': pose, 'intrinsics': intrinsics, 'scene-name': os.path.basename(scene_folder_path)} + if DEBUG: + break with open(os.path.join(dataset_path, 'cameras.json'), 'w') as outfile: json.dump(cameras, outfile, indent=4) @@ -69,7 +80,7 @@ def list_recursive(folderpath): camera_dataset_file = os.path.join(args.source, 'cameras.json') with open(camera_dataset_file, "r") as f: - cameras = json.load(f) + cameras = json.load(f) # same camera file as saved above dataset = {'labels':[]} max_images = args.max_images if args.max_images is not None else len(cameras) @@ -77,11 +88,16 @@ def list_recursive(folderpath): if (max_images is not None and i >= max_images): break pose = np.array(cameras[filename]['pose']) + # st() intrinsics = np.array(cameras[filename]['intrinsics']) label = np.concatenate([pose.reshape(-1), intrinsics.reshape(-1)]).tolist() image_path = os.path.join(args.source, filename) dataset["labels"].append([filename, label]) + # st() + # check cameras/dataset + print(os.path.join(args.source, 'dataset.json')) with open(os.path.join(args.source, 'dataset.json'), "w") as f: json.dump(dataset, f, indent=4) + diff --git a/dataset_preprocessing/shapenet_cars/run_me.py b/dataset_preprocessing/shapenet_cars/run_me.py index 894f33e8..18a06b27 100644 --- a/dataset_preprocessing/shapenet_cars/run_me.py +++ b/dataset_preprocessing/shapenet_cars/run_me.py @@ -17,19 +17,22 @@ if __name__ == '__main__': with tempfile.TemporaryDirectory() as working_dir: + # working_dir = '/tmp/tmphal02_sj' # /cars_train.zip + working_dir = '/home/xuyi/Data' + # print(working_dir) download_name = 'cars_train.zip' url = 'https://drive.google.com/uc?id=1bThUNtIHx4xEQyffVBSf82ABDDh2HlFn' - output_dataset_name = 'cars_128.zip' + output_dataset_name = 'cars_128_copy.zip' dir_path = os.path.dirname(os.path.realpath(__file__)) extracted_data_path = os.path.join(working_dir, os.path.splitext(download_name)[0]) print("Downloading data...") zipped_dataset = os.path.join(working_dir, download_name) - gdown.download(url, zipped_dataset, quiet=False) + # gdown.download(url, zipped_dataset, quiet=False) print("Unzipping downloaded data...") - shutil.unpack_archive(zipped_dataset, working_dir) + # shutil.unpack_archive(zipped_dataset, working_dir) print("Converting camera parameters...") cmd = f"python {os.path.join(dir_path, 'preprocess_shapenet_cameras.py')} --source={extracted_data_path}" diff --git a/eg3d/dataset_tool.py b/eg3d/dataset_tool.py index a400f770..5355a363 100644 --- a/eg3d/dataset_tool.py +++ b/eg3d/dataset_tool.py @@ -10,6 +10,9 @@ """Tool for creating ZIP/PNG based datasets.""" +from array import array +import csv +from email.policy import default import functools import gzip import io @@ -19,6 +22,7 @@ import re import sys import tarfile +from xmlrpc.client import boolean import zipfile from pathlib import Path from typing import Callable, Optional, Tuple, Union @@ -28,6 +32,10 @@ import PIL.Image from tqdm import tqdm +from ipdb import set_trace as st +import pandas as pd +import point_cloud_utils as pcu + #---------------------------------------------------------------------------- def error(msg): @@ -69,6 +77,11 @@ def is_image_ext(fname: Union[str, Path]) -> bool: def open_image_folder(source_dir, *, max_images: Optional[int]): input_images = [str(f) for f in sorted(Path(source_dir).rglob('*')) if is_image_ext(f) and os.path.isfile(f)] + + # filter out 'depth' and 'normal' + input_images = [ f for f in input_images if ('depth' not in f) and ('normal' not in f) and ('Image' not in f)] + # print(input_images) + # st() # Load labels. labels = {} @@ -76,19 +89,66 @@ def open_image_folder(source_dir, *, max_images: Optional[int]): if os.path.isfile(meta_fname): with open(meta_fname, 'r') as file: labels = json.load(file)['labels'] + # st() if labels is not None: + try: + pc_rel_paths = { x[0]: x[2] for x in labels } + # print(pc_rel_paths) + except: + print("No pointcloud input in dataset") + pc_rel_paths = {} labels = { x[0]: x[1] for x in labels } + else: labels = {} - - max_idx = maybe_min(len(input_images), max_images) + # print(labels) + + # max_idx = maybe_min(len(input_images), max_images) + max_idx = maybe_min(len(labels), max_images) def iterate_images(): for idx, fname in enumerate(input_images): arch_fname = os.path.relpath(fname, source_dir) arch_fname = arch_fname.replace('\\', '/') img = np.array(PIL.Image.open(fname)) - yield dict(img=img, label=labels.get(arch_fname)) + + if READ_POINTCLOUD: + + pc_rel = pc_rel_paths.get(arch_fname[:-4]) + if pc_rel != None: + pc_fname = os.path.join(source_dir, pc_rel) + pc_df = pd.read_csv(pc_fname) + pc_array = pc_df[['x','y','z','r','g','b','a', 'metallic','roughness']].values.astype(np.float32) + # st() + # reda pc csv + # save as np array + + if pc_array.shape[0] > NUM_POINTS: # poisson sampling + n_sample_poisson = NUM_POINTS + particle_pos = pc_array[:, :3] + poisson_idx = pcu.downsample_point_cloud_poisson_disk(particle_pos, num_samples=n_sample_poisson) + while poisson_idx.shape[0] < NUM_POINTS: + n_sample_poisson += 50 + poisson_idx = pcu.downsample_point_cloud_poisson_disk(particle_pos, num_samples=n_sample_poisson) + poisson_idx = poisson_idx[:NUM_POINTS] + # particle_pos = particle_pos[poisson_idx] + pc_array = pc_array[poisson_idx] + else: + continue + else: + pc_array=None + + + arch_fname = os.path.splitext(arch_fname)[0] + label_get = labels.get(arch_fname) + if label_get != None: + # st() + # print('fname, labels.get(fname)', arch_fname) + yield dict(img=img, label=labels.get(arch_fname), pc=pc_array) + else: + # print(arch_fname, ' is not written into json file yet') # in the abo case + # st() + pass if idx >= max_idx-1: break return max_idx, iterate_images() @@ -98,7 +158,7 @@ def iterate_images(): def open_image_zip(source, *, max_images: Optional[int]): with zipfile.ZipFile(source, mode='r') as z: input_images = [str(f) for f in sorted(z.namelist()) if is_image_ext(f)] - + st() # Load labels. labels = {} if 'dataset.json' in z.namelist(): @@ -117,6 +177,7 @@ def iterate_images(): with z.open(fname, 'r') as file: img = PIL.Image.open(file) # type: ignore img = np.array(img) + yield dict(img=img, label=labels.get(fname)) if idx >= max_idx-1: break @@ -266,6 +327,7 @@ def center_crop_wide(width, height, img): #---------------------------------------------------------------------------- def open_dataset(source, *, max_images: Optional[int]): + # st() if os.path.isdir(source): if source.rstrip('/').endswith('_lmdb'): return open_lmdb(source, max_images=max_images) @@ -294,6 +356,7 @@ def open_dest(dest: str) -> Tuple[str, Callable[[str, Union[bytes, str]], None], zf = zipfile.ZipFile(file=dest, mode='w', compression=zipfile.ZIP_STORED) def zip_write_bytes(fname: str, data: Union[bytes, str]): zf.writestr(fname, data) + return '', zip_write_bytes, zf.close else: # If the output folder already exists, check that is is @@ -306,6 +369,7 @@ def zip_write_bytes(fname: str, data: Union[bytes, str]): if os.path.isdir(dest) and len(os.listdir(dest)) != 0: error('--dest folder must be empty') os.makedirs(dest, exist_ok=True) + st() def folder_write_bytes(fname: str, data: Union[bytes, str]): os.makedirs(os.path.dirname(fname), exist_ok=True) @@ -324,13 +388,15 @@ def folder_write_bytes(fname: str, data: Union[bytes, str]): @click.option('--max-images', help='Output only up to `max-images` images', type=int, default=None) @click.option('--transform', help='Input crop/resize mode', type=click.Choice(['center-crop', 'center-crop-wide'])) @click.option('--resolution', help='Output resolution (e.g., \'512x512\')', metavar='WxH', type=parse_tuple) +@click.option('--read_pointcloud', help='whether pc.csv is in dataset)', type=boolean, is_flag=True, default=False) def convert_dataset( ctx: click.Context, source: str, dest: str, max_images: Optional[int], transform: Optional[str], - resolution: Optional[Tuple[int, int]] + resolution: Optional[Tuple[int, int]], + read_pointcloud:Optional[boolean] ): """Convert an image dataset into a dataset archive usable with StyleGAN2 ADA PyTorch. @@ -391,12 +457,19 @@ def convert_dataset( --transform=center-crop-wide --resolution=512x384 """ + if read_pointcloud: + global READ_POINTCLOUD + READ_POINTCLOUD = True + global NUM_POINTS + NUM_POINTS = 1024 + PIL.Image.init() # type: ignore if dest == '': ctx.fail('--dest output filename or directory must not be an empty string') num_files, input_iter = open_dataset(source, max_images=max_images) + archive_root_dir, save_bytes, close_dest = open_dest(dest) if resolution is None: resolution = (None, None) @@ -406,6 +479,7 @@ def convert_dataset( labels = [] for idx, image in tqdm(enumerate(input_iter), total=num_files): + # print(image) idx_str = f'{idx:08d}' archive_fname = f'{idx_str[:5]}/img{idx_str}.png' @@ -439,17 +513,42 @@ def convert_dataset( error(f'Image {archive_fname} attributes must be equal across all images of the dataset. Got:\n' + '\n'.join(err)) # Save the image as an uncompressed PNG. + WHITE_BKGD=True + if channels == 4 and WHITE_BKGD: + img = img[...,:3] * (img[...,-1:]/255) + (255 - img[...,-1:]) + img = img.astype(np.uint8) + channels = 3 + # im1 = img.save("geeks_white.png") + # st() img = PIL.Image.fromarray(img, { 1: 'L', 3: 'RGB', 4: 'RGBA'}[channels]) - if channels == 4: img = img.convert('RGB') + + if not WHITE_BKGD: + if channels == 4: img = img.convert('RGB') + # im1 = img.save("geeks_blk.png") + # st() + image_bits = io.BytesIO() img.save(image_bits, format='png', compress_level=0, optimize=False) save_bytes(os.path.join(archive_root_dir, archive_fname), image_bits.getbuffer()) + # print(archive_fname) labels.append([archive_fname, image['label']] if image['label'] is not None else None) + # Save pc.csv also using the same indexed archname + pc_array = image['pc'] + archive_fname_pc = archive_fname.replace('img', 'pc').replace('png', 'csv') + string_buffer = io.StringIO() + writer = csv.writer(string_buffer) + for row in pc_array: + writer.writerow(row) + save_bytes(os.path.join(archive_root_dir, archive_fname_pc), string_buffer.getvalue()) + print(archive_fname_pc) + metadata = { 'labels': labels if all(x is not None for x in labels) else None } + # print(metadata) save_bytes(os.path.join(archive_root_dir, 'dataset.json'), json.dumps(metadata)) + # print(os.path.join(archive_root_dir, 'dataset.json')) close_dest() #---------------------------------------------------------------------------- diff --git a/eg3d/environment.yml b/eg3d/environment.yml index 082bcaf5..9fcc1d0c 100644 --- a/eg3d/environment.yml +++ b/eg3d/environment.yml @@ -13,19 +13,20 @@ channels: - pytorch - nvidia dependencies: - - python >= 3.8 + - python>=3.9 - pip - numpy>=1.20 - - click>=8.0 - - pillow=8.3.1 - - scipy=1.7.1 - - pytorch=1.11.0 - - cudatoolkit=11.1 - - requests=2.26.0 - - tqdm=4.62.2 - - ninja=1.10.2 + - click + - pillow + - scipy + - nvidia::cudatoolkit=11.3 + - pytorch::pytorch=1.12.1=py3.9_cuda11.3_cudnn8.3.2_0 + - pytorch::torchvision=0.13.1=py39_cu113 + - requests + - tqdm + - ninja - matplotlib=3.4.2 - - imageio=2.9.0 + - imageio - pip: - imgui==1.3.0 - glfw==2.2.0 @@ -34,4 +35,13 @@ dependencies: - pyspng - psutil - mrcfile - - tensorboard \ No newline at end of file + - tensorboard + - ipdb + - pandas + - spconv + - https://data.pyg.org/whl/torch-1.12.0%2Bcu113/torch_scatter-2.0.9-cp39-cp39-linux_x86_64.whl + - ftfy + - regex + - dill + - git+https://github.com/openai/CLIP.git + - git+https://github.com/mapillary/inplace_abn \ No newline at end of file diff --git a/eg3d/gen_videos.py b/eg3d/gen_videos.py index de03d44c..323b5e08 100644 --- a/eg3d/gen_videos.py +++ b/eg3d/gen_videos.py @@ -12,6 +12,7 @@ import os import re + from typing import List, Optional, Tuple, Union import click @@ -27,6 +28,19 @@ from camera_utils import LookAtPoseSampler from torch_utils import misc +from ipdb import set_trace as st +import zipfile +import pandas as pd + +#---------------------------------------------------------------------------- +import ast +class PythonLiteralOption(click.Option): + + def type_cast_value(self, ctx, value): + try: + return ast.literal_eval(value) + except: + raise click.BadParameter(value) #---------------------------------------------------------------------------- def layout_grid(img, grid_w=None, grid_h=1, float_to_uint8=True, chw_to_hwc=True, to_numpy=True): @@ -91,12 +105,29 @@ def gen_interp_video(G, mp4: str, seeds, shuffle_seed=None, w_frames=60*4, kind= camera_lookat_point = torch.tensor(G.rendering_kwargs['avg_camera_pivot'], device=device) zs = torch.from_numpy(np.stack([np.random.RandomState(seed).randn(G.z_dim) for seed in all_seeds])).to(device) cam2world_pose = LookAtPoseSampler.sample(3.14/2, 3.14/2, camera_lookat_point, radius=G.rendering_kwargs['avg_camera_radius'], device=device) - focal_length = 4.2647 if cfg != 'Shapenet' else 1.7074 # shapenet has higher FOV + + # focal_length = 4.2647 if (cfg != 'Shapenet' or cfg != 'ABO') else 1.7074 # shapenet has higher FOV + if (cfg != 'Shapenet' and cfg != 'ABO'): + focal_length = 4.2647 + elif cfg == 'Shapenet': + focal_length = 1.7074 + elif cfg == 'ABO': + focal_length = 0.3889 + else: + print("Not supported dataset type") + print("Focal length: ", focal_length) intrinsics = torch.tensor([[focal_length, 0, 0.5], [0, focal_length, 0.5], [0, 0, 1]], device=device) c = torch.cat([cam2world_pose.reshape(-1, 16), intrinsics.reshape(-1, 9)], 1) c = c.repeat(len(zs), 1) + ws = G.mapping(z=zs, c=c, truncation_psi=psi, truncation_cutoff=truncation_cutoff) - _ = G.synthesis(ws[:1], c[:1]) # warm up + + if cfg == 'ABO': + # st() + _ = G.synthesis(ws[:1], c[:1], pc=PC_FILES[:1]) + else: + # st() + _ = G.synthesis(ws[:1], c[:1]) # warm up ws = ws.reshape(grid_h, grid_w, num_keyframes, *ws.shape[1:]) # Interpolation. @@ -128,13 +159,14 @@ def gen_interp_video(G, mp4: str, seeds, shuffle_seed=None, w_frames=60*4, kind= cam2world_pose = LookAtPoseSampler.sample(3.14/2 + yaw_range * np.sin(2 * 3.14 * frame_idx / (num_keyframes * w_frames)), 3.14/2 -0.05 + pitch_range * np.cos(2 * 3.14 * frame_idx / (num_keyframes * w_frames)), camera_lookat_point, radius=G.rendering_kwargs['avg_camera_radius'], device=device) + all_poses.append(cam2world_pose.squeeze().cpu().numpy()) - focal_length = 4.2647 if cfg != 'Shapenet' else 1.7074 # shapenet has higher FOV + focal_length = 4.2647 if (cfg != 'Shapenet' or cfg != 'ABO') else 1.7074 # shapenet has higher FOV intrinsics = torch.tensor([[focal_length, 0, 0.5], [0, focal_length, 0.5], [0, 0, 1]], device=device) c = torch.cat([cam2world_pose.reshape(-1, 16), intrinsics.reshape(-1, 9)], 1) interp = grid[yi][xi] - w = torch.from_numpy(interp(frame_idx / w_frames)).to(device) + w = torch.from_numpy(interp(frame_idx / w_frames).astype(np.float32)).to(device) entangle = 'camera' if entangle == 'conditioning': @@ -145,7 +177,10 @@ def gen_interp_video(G, mp4: str, seeds, shuffle_seed=None, w_frames=60*4, kind= w_c = G.mapping(z=zs[0:1], c=c[0:1], truncation_psi=psi, truncation_cutoff=truncation_cutoff) img = G.synthesis(ws=w_c, c=c_forward, noise_mode='const')[image_mode][0] elif entangle == 'camera': - img = G.synthesis(ws=w.unsqueeze(0), c=c[0:1], noise_mode='const')[image_mode][0] + if cfg == 'ABO': + pass + # st() + img = G.synthesis(ws=w.unsqueeze(0), c=c[0:1], pc=PC_FILES[0:1], noise_mode='const')[image_mode][0] elif entangle == 'both': w_c = G.mapping(z=zs[0:1], c=c[0:1], truncation_psi=psi, truncation_cutoff=truncation_cutoff) img = G.synthesis(ws=w_c, c=c[0:1], noise_mode='const')[image_mode][0] @@ -200,6 +235,7 @@ def gen_interp_video(G, mp4: str, seeds, shuffle_seed=None, w_frames=60*4, kind= video_out.close() all_poses = np.stack(all_poses) + if gen_shapes: print(all_poses.shape) with open(mp4.replace('.mp4', '_trajectory.npy'), 'wb') as f: @@ -249,12 +285,14 @@ def parse_tuple(s: Union[str, Tuple[int,int]]) -> Tuple[int, int]: @click.option('--trunc-cutoff', 'truncation_cutoff', type=int, help='Truncation cutoff', default=14, show_default=True) @click.option('--outdir', help='Output directory', type=str, required=True, metavar='DIR') @click.option('--reload_modules', help='Overload persistent modules?', type=bool, required=False, metavar='BOOL', default=False, show_default=True) -@click.option('--cfg', help='Config', type=click.Choice(['FFHQ', 'AFHQ', 'Shapenet']), required=False, metavar='STR', default='FFHQ', show_default=True) +@click.option('--cfg', help='Config', type=click.Choice(['FFHQ', 'AFHQ', 'Shapenet', 'ABO']), required=False, metavar='STR', default='FFHQ', show_default=True) @click.option('--image_mode', help='Image mode', type=click.Choice(['image', 'image_depth', 'image_raw']), required=False, metavar='STR', default='image', show_default=True) @click.option('--sample_mult', 'sampling_multiplier', type=float, help='Multiplier for depth sampling in volume rendering', default=2, show_default=True) @click.option('--nrr', type=int, help='Neural rendering resolution override', default=None, show_default=True) @click.option('--shapes', type=bool, help='Gen shapes for shape interpolation', default=False, show_default=True) @click.option('--interpolate', type=bool, help='Interpolate between seeds', default=True, show_default=True) +@click.option('--pointcloud_files', cls=PythonLiteralOption, default=[]) +@click.option('--data_zip', help='Dataset in zip format', type=str, required=False, metavar='DIR') def generate_images( network_pkl: str, @@ -268,6 +306,8 @@ def generate_images( outdir: str, reload_modules: bool, cfg: str, + pointcloud_files: List[str], + data_zip: str, image_mode: str, sampling_multiplier: float, nrr: Optional[int], @@ -314,6 +354,52 @@ def generate_images( if truncation_psi == 1.0: truncation_cutoff = 14 # no truncation so doesn't matter where we cutoff + + # st() + global PC_FILES + + def _file_ext(fname): + return os.path.splitext(fname)[1].lower() + + def _get_zipfile(): + # assert self._type == 'zip' + _zipfile = zipfile.ZipFile(data_zip) + return _zipfile + + + def _load_raw_pointcloud(raw_idx): + fname = _pc_fnames[raw_idx] + + with _get_zipfile().open(fname, 'r') as f: + df = pd.read_csv(f, header=None) + pc_array = df.values.astype(np.float32) + return pc_array + + + def _load_raw_pointcloud_by_name(f): + # fname = _pc_fnames[raw_idx] + + # with _get_zipfile().open(fname, 'r') as f: + pc_df = pd.read_csv(f) + pc_array = pc_df[['x','y','z','r','g','b','a', 'metallic','roughness']].values.astype(np.float32) + # pc_array = df.values.astype(np.float32) + return pc_array + + + if len(pointcloud_files) !=0: + # PC_FILES = pointcloud_files + PC_FILES = torch.tensor(np.asarray([_load_raw_pointcloud_by_name(i) for i in pointcloud_files]), device=device) + PC_FILES = PC_FILES.repeat(4,1,1) + else: + print("use predefined pointcloud") + _all_fnames = set(_get_zipfile().namelist()) + _pc_fnames = sorted(fname for fname in _all_fnames if _file_ext(fname) == '.csv') + # st() + indices = [205,307, 0,102] + PC_FILES = torch.tensor(np.asarray([_load_raw_pointcloud(i) for i in indices]), device=device) # B, 1024, 9 + + + if interpolate: output = os.path.join(outdir, 'interpolation.mp4') gen_interp_video(G=G, mp4=output, bitrate='10M', grid_dims=grid, num_keyframes=num_keyframes, w_frames=w_frames, seeds=seeds, shuffle_seed=shuffle_seed, psi=truncation_psi, truncation_cutoff=truncation_cutoff, cfg=cfg, image_mode=image_mode, gen_shapes=shapes) @@ -326,6 +412,8 @@ def generate_images( #---------------------------------------------------------------------------- if __name__ == "__main__": + # global pointcloud_files + generate_images() # pylint: disable=no-value-for-parameter #---------------------------------------------------------------------------- diff --git a/eg3d/gen_videos_c2w.py b/eg3d/gen_videos_c2w.py new file mode 100644 index 00000000..f97b73f9 --- /dev/null +++ b/eg3d/gen_videos_c2w.py @@ -0,0 +1,466 @@ +# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +"""Generate lerp videos using pretrained network pickle.""" + +import os +import re + +from typing import List, Optional, Tuple, Union + +import click +import dnnlib +import imageio +import numpy as np +import scipy.interpolate +import torch +from tqdm import tqdm +import mrcfile + +import legacy + +from camera_utils import LookAtPoseSampler +from torch_utils import misc +from ipdb import set_trace as st +import zipfile +import pandas as pd + +#---------------------------------------------------------------------------- +import ast +class PythonLiteralOption(click.Option): + + def type_cast_value(self, ctx, value): + try: + return ast.literal_eval(value) + except: + raise click.BadParameter(value) +#---------------------------------------------------------------------------- + +def layout_grid(img, grid_w=None, grid_h=1, float_to_uint8=True, chw_to_hwc=True, to_numpy=True): + batch_size, channels, img_h, img_w = img.shape + if grid_w is None: + grid_w = batch_size // grid_h + assert batch_size == grid_w * grid_h + if float_to_uint8: + img = (img * 127.5 + 128).clamp(0, 255).to(torch.uint8) + img = img.reshape(grid_h, grid_w, channels, img_h, img_w) + img = img.permute(2, 0, 3, 1, 4) + img = img.reshape(channels, grid_h * img_h, grid_w * img_w) + if chw_to_hwc: + img = img.permute(1, 2, 0) + if to_numpy: + img = img.cpu().numpy() + return img + +def create_samples(N=256, voxel_origin=[0, 0, 0], cube_length=2.0): + # NOTE: the voxel_origin is actually the (bottom, left, down) corner, not the middle + voxel_origin = np.array(voxel_origin) - cube_length/2 + voxel_size = cube_length / (N - 1) + + overall_index = torch.arange(0, N ** 3, 1, out=torch.LongTensor()) + samples = torch.zeros(N ** 3, 3) + + # transform first 3 columns + # to be the x, y, z index + samples[:, 2] = overall_index % N + samples[:, 1] = (overall_index.float() / N) % N + samples[:, 0] = ((overall_index.float() / N) / N) % N + + # transform first 3 columns + # to be the x, y, z coordinate + samples[:, 0] = (samples[:, 0] * voxel_size) + voxel_origin[2] + samples[:, 1] = (samples[:, 1] * voxel_size) + voxel_origin[1] + samples[:, 2] = (samples[:, 2] * voxel_size) + voxel_origin[0] + + num_samples = N ** 3 + + return samples.unsqueeze(0), voxel_origin, voxel_size + +#---------------------------------------------------------------------------- + +def gen_interp_video(G, mp4: str, seeds, shuffle_seed=None, w_frames=60*4, kind='cubic', grid_dims=(1,1), num_keyframes=None, wraps=2, psi=1, truncation_cutoff=14, cfg='FFHQ', image_mode='image', gen_shapes=False, device=torch.device('cuda'), **video_kwargs): + grid_w = grid_dims[0] + grid_h = grid_dims[1] + + if num_keyframes is None: + if len(seeds) % (grid_w*grid_h) != 0: + raise ValueError('Number of input seeds must be divisible by grid W*H') + num_keyframes = len(seeds) // (grid_w*grid_h) + + all_seeds = np.zeros(num_keyframes*grid_h*grid_w, dtype=np.int64) + for idx in range(num_keyframes*grid_h*grid_w): + all_seeds[idx] = seeds[idx % len(seeds)] + + if shuffle_seed is not None: + rng = np.random.RandomState(seed=shuffle_seed) + rng.shuffle(all_seeds) + + camera_lookat_point = torch.tensor(G.rendering_kwargs['avg_camera_pivot'], device=device) + zs = torch.from_numpy(np.stack([np.random.RandomState(seed).randn(G.z_dim) for seed in all_seeds])).to(device) + cam2world_pose = LookAtPoseSampler.sample(3.14/2, 3.14/2, camera_lookat_point, radius=G.rendering_kwargs['avg_camera_radius'], device=device) + + # focal_length = 4.2647 if (cfg != 'Shapenet' or cfg != 'ABO') else 1.7074 # shapenet has higher FOV + if (cfg != 'Shapenet' and cfg != 'ABO'): + focal_length = 4.2647 + elif cfg == 'Shapenet': + focal_length = 1.7074 + elif cfg == 'ABO': + focal_length = 1.3889 + else: + print("Not supported dataset type") + print("Focal length: ", focal_length) + intrinsics = torch.tensor([[focal_length, 0, 0.5], [0, focal_length, 0.5], [0, 0, 1]], device=device) + c = torch.cat([cam2world_pose.reshape(-1, 16), intrinsics.reshape(-1, 9)], 1) + c = c.repeat(len(zs), 1) + + ws = G.mapping(z=zs, c=c, truncation_psi=psi, truncation_cutoff=truncation_cutoff) + + if cfg == 'ABO': + # st() + _ = G.synthesis(ws[:1], c[:1], pc=PC_FILES[:1]) + else: + # st() + _ = G.synthesis(ws[:1], c[:1]) # warm up + ws = ws.reshape(grid_h, grid_w, num_keyframes, *ws.shape[1:]) + + # Interpolation. + grid = [] + for yi in range(grid_h): + row = [] + for xi in range(grid_w): + x = np.arange(-num_keyframes * wraps, num_keyframes * (wraps + 1)) + y = np.tile(ws[yi][xi].cpu().numpy(), [wraps * 2 + 1, 1, 1]) + interp = scipy.interpolate.interp1d(x, y, kind=kind, axis=0) + row.append(interp) + grid.append(row) + + # Render video. + max_batch = 10000000 + voxel_resolution = 512 + video_out = imageio.get_writer(mp4, mode='I', fps=60, codec='libx264', **video_kwargs) + + if gen_shapes: + outdir = 'interpolation_{}_{}/'.format(all_seeds[0], all_seeds[1]) + os.makedirs(outdir, exist_ok=True) + + + + + all_poses = [] + for frame_idx in tqdm(range(num_keyframes * w_frames)): + imgs = [] + for yi in range(grid_h): + for xi in range(grid_w): + pitch_range = 0.25 + yaw_range = 0.35 + try: + c_idx = frame_idx%len(PREDEFINED_POSES) + # print(c_idx) + cam2world_pose = PREDEFINED_POSES[c_idx:c_idx+1] + except: + st() + cam2world_pose = LookAtPoseSampler.sample(3.14/2 + yaw_range * np.sin(2 * 3.14 * frame_idx / (num_keyframes * w_frames)), + 3.14/2 -0.05 + pitch_range * np.cos(2 * 3.14 * frame_idx / (num_keyframes * w_frames)), + camera_lookat_point, radius=G.rendering_kwargs['avg_camera_radius'], device=device) + + all_poses.append(cam2world_pose.squeeze().cpu().numpy()) + # focal_length = 4.2647 if (cfg != 'Shapenet' and cfg != 'ABO') else 1.7074 # shapenet has higher FOV + if (cfg != 'Shapenet' and cfg != 'ABO'): + focal_length = 4.2647 + elif cfg == 'Shapenet': + focal_length = 1.7074 + elif cfg == 'ABO': + focal_length = 1.3889 + else: + print("Not supported dataset type") + intrinsics = torch.tensor([[focal_length, 0, 0.5], [0, focal_length, 0.5], [0, 0, 1]], device=device) + c = torch.cat([cam2world_pose.reshape(-1, 16), intrinsics.reshape(-1, 9)], 1) + # st() + + interp = grid[yi][xi] + w = torch.from_numpy(interp(frame_idx / w_frames).astype(np.float32)).to(device) + + entangle = 'camera' + if entangle == 'conditioning': + c_forward = torch.cat([LookAtPoseSampler.sample(3.14/2, + 3.14/2, + camera_lookat_point, + radius=G.rendering_kwargs['avg_camera_radius'], device=device).reshape(-1, 16), intrinsics.reshape(-1, 9)], 1) + w_c = G.mapping(z=zs[0:1], c=c[0:1], truncation_psi=psi, truncation_cutoff=truncation_cutoff) + img = G.synthesis(ws=w_c, c=c_forward, noise_mode='const')[image_mode][0] + elif entangle == 'camera': + if cfg == 'ABO': + pass + # st() + img = G.synthesis(ws=w.unsqueeze(0), c=c[0:1], pc=PC_FILES[0:1], noise_mode='const')[image_mode][0] + elif entangle == 'both': + w_c = G.mapping(z=zs[0:1], c=c[0:1], truncation_psi=psi, truncation_cutoff=truncation_cutoff) + img = G.synthesis(ws=w_c, c=c[0:1], noise_mode='const')[image_mode][0] + + if image_mode == 'image_depth': + img = -img + img = (img - img.min()) / (img.max() - img.min()) * 2 - 1 + + imgs.append(img) + + if gen_shapes: + # generate shapes + print('Generating shape for frame %d / %d ...' % (frame_idx, num_keyframes * w_frames)) + + samples, voxel_origin, voxel_size = create_samples(N=voxel_resolution, voxel_origin=[0, 0, 0], cube_length=G.rendering_kwargs['box_warp']) + samples = samples.to(device) + sigmas = torch.zeros((samples.shape[0], samples.shape[1], 1), device=device) + transformed_ray_directions_expanded = torch.zeros((samples.shape[0], max_batch, 3), device=device) + transformed_ray_directions_expanded[..., -1] = -1 + + head = 0 + with tqdm(total = samples.shape[1]) as pbar: + with torch.no_grad(): + while head < samples.shape[1]: + torch.manual_seed(0) + sigma = G.sample_mixed(samples[:, head:head+max_batch], transformed_ray_directions_expanded[:, :samples.shape[1]-head], w.unsqueeze(0), truncation_psi=psi, noise_mode='const')['sigma'] + sigmas[:, head:head+max_batch] = sigma + head += max_batch + pbar.update(max_batch) + + sigmas = sigmas.reshape((voxel_resolution, voxel_resolution, voxel_resolution)).cpu().numpy() + sigmas = np.flip(sigmas, 0) + + pad = int(30 * voxel_resolution / 256) + pad_top = int(38 * voxel_resolution / 256) + sigmas[:pad] = 0 + sigmas[-pad:] = 0 + sigmas[:, :pad] = 0 + sigmas[:, -pad_top:] = 0 + sigmas[:, :, :pad] = 0 + sigmas[:, :, -pad:] = 0 + + output_ply = True + if output_ply: + from shape_utils import convert_sdf_samples_to_ply + convert_sdf_samples_to_ply(np.transpose(sigmas, (2, 1, 0)), [0, 0, 0], 1, os.path.join(outdir, f'{frame_idx:04d}_shape.ply'), level=10) + else: # output mrc + with mrcfile.new_mmap(outdir + f'{frame_idx:04d}_shape.mrc', overwrite=True, shape=sigmas.shape, mrc_mode=2) as mrc: + mrc.data[:] = sigmas + + video_out.append_data(layout_grid(torch.stack(imgs), grid_w=grid_w, grid_h=grid_h)) + video_out.close() + all_poses = np.stack(all_poses) + + + if gen_shapes: + print(all_poses.shape) + with open(mp4.replace('.mp4', '_trajectory.npy'), 'wb') as f: + np.save(f, all_poses) + +#---------------------------------------------------------------------------- + +def parse_range(s: Union[str, List[int]]) -> List[int]: + '''Parse a comma separated list of numbers or ranges and return a list of ints. + + Example: '1,2,5-10' returns [1, 2, 5, 6, 7] + ''' + if isinstance(s, list): return s + ranges = [] + range_re = re.compile(r'^(\d+)-(\d+)$') + for p in s.split(','): + if m := range_re.match(p): + ranges.extend(range(int(m.group(1)), int(m.group(2))+1)) + else: + ranges.append(int(p)) + return ranges + +#---------------------------------------------------------------------------- + +def parse_tuple(s: Union[str, Tuple[int,int]]) -> Tuple[int, int]: + '''Parse a 'M,N' or 'MxN' integer tuple. + + Example: + '4x2' returns (4,2) + '0,1' returns (0,1) + ''' + if isinstance(s, tuple): return s + if m := re.match(r'^(\d+)[x,](\d+)$', s): + return (int(m.group(1)), int(m.group(2))) + raise ValueError(f'cannot parse tuple {s}') + +#---------------------------------------------------------------------------- + +@click.command() +@click.option('--network', 'network_pkl', help='Network pickle filename', required=True) +@click.option('--seeds', type=parse_range, help='List of random seeds', required=True) +@click.option('--shuffle-seed', type=int, help='Random seed to use for shuffling seed order', default=None) +@click.option('--grid', type=parse_tuple, help='Grid width/height, e.g. \'4x3\' (default: 1x1)', default=(1,1)) +@click.option('--num-keyframes', type=int, help='Number of seeds to interpolate through. If not specified, determine based on the length of the seeds array given by --seeds.', default=None) +@click.option('--w-frames', type=int, help='Number of frames to interpolate between latents', default=120) +@click.option('--trunc', 'truncation_psi', type=float, help='Truncation psi', default=1, show_default=True) +@click.option('--trunc-cutoff', 'truncation_cutoff', type=int, help='Truncation cutoff', default=14, show_default=True) +@click.option('--outdir', help='Output directory', type=str, required=True, metavar='DIR') +@click.option('--reload_modules', help='Overload persistent modules?', type=bool, required=False, metavar='BOOL', default=False, show_default=True) +@click.option('--cfg', help='Config', type=click.Choice(['FFHQ', 'AFHQ', 'Shapenet', 'ABO']), required=False, metavar='STR', default='FFHQ', show_default=True) +@click.option('--image_mode', help='Image mode', type=click.Choice(['image', 'image_depth', 'image_raw']), required=False, metavar='STR', default='image', show_default=True) +@click.option('--sample_mult', 'sampling_multiplier', type=float, help='Multiplier for depth sampling in volume rendering', default=2, show_default=True) +@click.option('--nrr', type=int, help='Neural rendering resolution override', default=None, show_default=True) +@click.option('--shapes', type=bool, help='Gen shapes for shape interpolation', default=False, show_default=True) +@click.option('--interpolate', type=bool, help='Interpolate between seeds', default=True, show_default=True) +@click.option('--pointcloud_files', cls=PythonLiteralOption, default=[]) +@click.option('--data_zip', help='Dataset in zip format', type=str, required=False, metavar='DIR') +@click.option('--pose_file', help='predefined c2ws for abo dataset', type=str, required=False, metavar='DIR') + +def generate_images( + network_pkl: str, + seeds: List[int], + shuffle_seed: Optional[int], + truncation_psi: float, + truncation_cutoff: int, + grid: Tuple[int,int], + num_keyframes: Optional[int], + w_frames: int, + outdir: str, + reload_modules: bool, + cfg: str, + pointcloud_files: List[str], + pose_file: str, + data_zip: str, + image_mode: str, + sampling_multiplier: float, + nrr: Optional[int], + shapes: bool, + interpolate: bool, +): + """Render a latent vector interpolation video. + + Examples: + + \b + # Render a 4x2 grid of interpolations for seeds 0 through 31. + python gen_video.py --output=lerp.mp4 --trunc=1 --seeds=0-31 --grid=4x2 \\ + --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-afhqv2-512x512.pkl + + Animation length and seed keyframes: + + The animation length is either determined based on the --seeds value or explicitly + specified using the --num-keyframes option. + + When num keyframes is specified with --num-keyframes, the output video length + will be 'num_keyframes*w_frames' frames. + + If --num-keyframes is not specified, the number of seeds given with + --seeds must be divisible by grid size W*H (--grid). In this case the + output video length will be '# seeds/(w*h)*w_frames' frames. + """ + + if not os.path.exists(outdir): + os.makedirs(outdir, exist_ok=True) + + print('Loading networks from "%s"...' % network_pkl) + device = torch.device('cuda') + with dnnlib.util.open_url(network_pkl) as f: + G = legacy.load_network_pkl(f)['G_ema'].to(device) # type: ignore + + + G.rendering_kwargs['depth_resolution'] = int(G.rendering_kwargs['depth_resolution'] * sampling_multiplier) + G.rendering_kwargs['depth_resolution_importance'] = int(G.rendering_kwargs['depth_resolution_importance'] * sampling_multiplier) + if nrr is not None: G.neural_rendering_resolution = nrr + + if truncation_cutoff == 0: + truncation_psi = 1.0 # truncation cutoff of 0 means no truncation anyways + if truncation_psi == 1.0: + truncation_cutoff = 14 # no truncation so doesn't matter where we cutoff + + + # st() + ################################################ + global PC_FILES + + def _file_ext(fname): + return os.path.splitext(fname)[1].lower() + + def _get_zipfile(): + # assert self._type == 'zip' + _zipfile = zipfile.ZipFile(data_zip) + return _zipfile + + def _load_raw_pointcloud(raw_idx): + fname = _pc_fnames[raw_idx] + + with _get_zipfile().open(fname, 'r') as f: + df = pd.read_csv(f, header=None) + pc_array = df.values.astype(np.float32) + return pc_array + + def _load_raw_pointcloud_by_name(f): + # fname = _pc_fnames[raw_idx] + + # with _get_zipfile().open(fname, 'r') as f: + pc_df = pd.read_csv(f) + pc_array = pc_df[['x','y','z','r','g','b','a', 'metallic','roughness']].values.astype(np.float32) + # pc_array = df.values.astype(np.float32) + return pc_array + + + if len(pointcloud_files) !=0: + # PC_FILES = pointcloud_files + PC_FILES = torch.tensor(np.asarray([_load_raw_pointcloud_by_name(i) for i in pointcloud_files]), device=device) + PC_FILES = PC_FILES.repeat(4,1,1) + else: + print("use predefined pointcloud") + _all_fnames = set(_get_zipfile().namelist()) + _pc_fnames = sorted(fname for fname in _all_fnames if _file_ext(fname) == '.csv') + # st() + indices = [205,307, 0,102] + PC_FILES = torch.tensor(np.asarray([_load_raw_pointcloud(i) for i in indices]), device=device) # B, 1024, 9 + + ################################################ + global PREDEFINED_POSES + + if pose_file is not None: + import json + + with open(pose_file, 'r') as f: + meta = json.load(f) + all_poses=[] + blender2opencv = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]]) + for frame in meta ['frames']: + # rgb_path = frame['file_path'] + # relative_path = os.path.relpath(rgb_path, dataset_path) + # print(relative_path) + + # intrinsics = intrinsic_for_all + pose = (np.array(frame['transform_matrix'])@blender2opencv) + # print(len(pose)) + + # cameras[relative_path] = {'pose': pose, 'intrinsics': intrinsics, 'scene-name': os.path.basename(scene_folder_path),\ + # 'pc_csv':pc_relative_path} + all_poses.append(pose) + + PREDEFINED_POSES = torch.tensor(np.stack(all_poses), device=device) + + + + + + if interpolate: + output = os.path.join(outdir, 'interpolation.mp4') + gen_interp_video(G=G, mp4=output, bitrate='10M', grid_dims=grid, num_keyframes=num_keyframes, w_frames=w_frames, seeds=seeds, shuffle_seed=shuffle_seed, psi=truncation_psi, truncation_cutoff=truncation_cutoff, cfg=cfg, image_mode=image_mode, gen_shapes=shapes) + else: + for seed in seeds: + output = os.path.join(outdir, f'{seed}.mp4') + seeds_ = [seed] + gen_interp_video(G=G, mp4=output, bitrate='10M', grid_dims=grid, num_keyframes=num_keyframes, w_frames=w_frames, seeds=seeds_, shuffle_seed=shuffle_seed, psi=truncation_psi, truncation_cutoff=truncation_cutoff, cfg=cfg, image_mode=image_mode) + +#---------------------------------------------------------------------------- + +if __name__ == "__main__": + # global pointcloud_files + + generate_images() # pylint: disable=no-value-for-parameter + +#---------------------------------------------------------------------------- diff --git a/eg3d/metrics/metric_utils.py b/eg3d/metrics/metric_utils.py index 212cb7d3..86e7b64b 100644 --- a/eg3d/metrics/metric_utils.py +++ b/eg3d/metrics/metric_utils.py @@ -20,6 +20,9 @@ import torch import dnnlib +from ipdb import set_trace as st +from training.volume import VolumeGenerator + #---------------------------------------------------------------------------- class MetricOptions: @@ -57,16 +60,27 @@ def get_feature_detector(url, device=torch.device('cpu'), num_gpus=1, rank=0, ve #---------------------------------------------------------------------------- def iterate_random_labels(opts, batch_size): + if opts.G.c_dim == 0: c = torch.zeros([batch_size, opts.G.c_dim], device=opts.device) while True: yield c else: dataset = dnnlib.util.construct_class_by_name(**opts.dataset_kwargs) - while True: - c = [dataset.get_label(np.random.randint(len(dataset))) for _i in range(batch_size)] - c = torch.from_numpy(np.stack(c)).pin_memory().to(opts.device) - yield c + + if isinstance(opts.G, VolumeGenerator): + # to add get label and pc at the same time + while True: + c_and_pc = [dataset.get_label_and_pc(np.random.randint(len(dataset))) for _i in range(batch_size)] + c, pc = list(map(list, zip(*c_and_pc))) + c = torch.from_numpy(np.stack(c)).pin_memory().to(opts.device) + pc = torch.from_numpy(np.stack(pc)).pin_memory().to(opts.device) + yield c, pc + else: + while True: + c = [dataset.get_label(np.random.randint(len(dataset))) for _i in range(batch_size)] + c = torch.from_numpy(np.stack(c)).pin_memory().to(opts.device) + yield c #---------------------------------------------------------------------------- @@ -230,7 +244,7 @@ def compute_feature_stats_for_dataset(opts, detector_url, detector_kwargs, rel_l # Main loop. item_subset = [(i * opts.num_gpus + opts.rank) % num_items for i in range((num_items - 1) // opts.num_gpus + 1)] - for images, _labels in torch.utils.data.DataLoader(dataset=dataset, sampler=item_subset, batch_size=batch_size, **data_loader_kwargs): + for images, _labels, _pc in torch.utils.data.DataLoader(dataset=dataset, sampler=item_subset, batch_size=batch_size, **data_loader_kwargs): if images.shape[1] == 1: images = images.repeat([1, 3, 1, 1]) features = detector(images.to(opts.device), **detector_kwargs) @@ -253,6 +267,7 @@ def compute_feature_stats_for_generator(opts, detector_url, detector_kwargs, rel assert batch_size % batch_gen == 0 # Setup generator and labels. + G = copy.deepcopy(opts.G).eval().requires_grad_(False).to(opts.device) c_iter = iterate_random_labels(opts=opts, batch_size=batch_gen) @@ -265,11 +280,19 @@ def compute_feature_stats_for_generator(opts, detector_url, detector_kwargs, rel # Main loop. while not stats.is_full(): images = [] - for _i in range(batch_size // batch_gen): - z = torch.randn([batch_gen, G.z_dim], device=opts.device) - img = G(z=z, c=next(c_iter), **opts.G_kwargs)['image'] - img = (img * 127.5 + 128).clamp(0, 255).to(torch.uint8) - images.append(img) + if isinstance(opts.G, VolumeGenerator): # c_iter generate both c and pc + for _i in range(batch_size // batch_gen): + z = torch.randn([batch_gen, G.z_dim], device=opts.device) + c, pc = next(c_iter) + img = G(z=z, c=c, pc=pc, **opts.G_kwargs)['image'] + img = (img * 127.5 + 128).clamp(0, 255).to(torch.uint8) + images.append(img) + else: + for _i in range(batch_size // batch_gen): + z = torch.randn([batch_gen, G.z_dim], device=opts.device) + img = G(z=z, c=next(c_iter), **opts.G_kwargs)['image'] + img = (img * 127.5 + 128).clamp(0, 255).to(torch.uint8) + images.append(img) images = torch.cat(images) if images.shape[1] == 1: images = images.repeat([1, 3, 1, 1]) diff --git a/eg3d/scripts/gen_video.sh b/eg3d/scripts/gen_video.sh new file mode 100644 index 00000000..b4a999a8 --- /dev/null +++ b/eg3d/scripts/gen_video.sh @@ -0,0 +1,19 @@ +python gen_videos_c2w.py --outdir=out --trunc=0.7 --seeds=0-3 --grid=2x2 \ + --cfg ABO --pointcloud_files "['/home/xuyi/Data/renderer/output_abo/B07QJJKZL2/sample/pc.csv']" \ + --pose_file /home/xuyi/Data/renderer/output_abo/B07DBJX741/render/transforms.json \ + --network=/home/xuyi/Repo/eg3d/eg3d/pretrained_models/network-snapshot-synunet-000800.pkl + +# python gen_videos.py --outdir=out --trunc=0.7 --seeds=0-3 --grid=2x2 \ +# --cfg ABO --pointcloud_files "['/home/xuyi/Data/renderer/output_abo/B07JPGPBL2/sample/pc.csv']" \ +# --network=/home/xuyi/Repo/eg3d/eg3d/pretrained_models/network-snapshot-synunet-000600.pkl +# # --network=/home/xuyi/Repo/eg3d/eg3d/pretrained_models/network-snapshot-original-abo-001600.pkl +# # + +# # 'B07B4MHTG1', 'B07B4MF6P2', 'B07RTZ54B1', 'B075X4F3Z2', 'B07JL5QBC2', 'B075YPKYM1', 'B073NZGLT1', 'B07DYK2Y61' +# B07JXF7251: 花枕头 +# B07JPGPBL2：红椅子 +# B07QFP4M23: brown sofa +# B07QJJKZL2: blue blanket + +#### not-so-good cases +# B07QHKQMY4L red cup diff --git a/eg3d/scripts/jialin_depth_vis.sh b/eg3d/scripts/jialin_depth_vis.sh new file mode 100644 index 00000000..d871825c --- /dev/null +++ b/eg3d/scripts/jialin_depth_vis.sh @@ -0,0 +1,11 @@ +# ABO +##-------- common settings --------------------- +CUDA_VISIBLE_DEVICES=0 +GPUS=1 +BASE_DIR=/home/xuyi/Repo/eg3d +DATA=${BASE_DIR}/dataset_preprocessing/abo/abo_128_completed_white.zip +BATCH_SIZE=1 +python train.py --outdir=${BASE_DIR}/try-runs --cfg=abo_dataset --data=${DATA} \ + --gpus=${GPUS} --batch=${BATCH_SIZE} --gamma=0.3 \ + --backbone volume --decoder_dim 8 \ + --resume=/home/xuyi/Repo/eg3d/dataset_preprocessing/abo/network-snapshot-000600.pkl \ No newline at end of file diff --git a/eg3d/scripts/shapenet_car_finetune.sh b/eg3d/scripts/shapenet_car_finetune.sh new file mode 100644 index 00000000..ef8fa93d --- /dev/null +++ b/eg3d/scripts/shapenet_car_finetune.sh @@ -0,0 +1,43 @@ +# Train with Shapenet finetune, using 1 GPUs. +##-------- common settings --------------------- +CUDA_VISIBLE_DEVICES=0 +GPUS=2 +BATCH_SIZE=4 +BASE_DIR=/home/xuyi/Repo/eg3d + +# ##-------- abo/shapenet with triplane ----------- +# # DATA=${BASE_DIR}/dataset_preprocessing/shapenet_cars/cars_128_copy.zip +# DATA=${BASE_DIR}/dataset_preprocessing/abo/abo_128_copy.zip +# PRETRAINED_MODEL=${BASE_DIR}/eg3d/pretrained_models/shapenetcars128-64.pkl + +# python train.py --outdir=${BASE_DIR}/try-runs --cfg=shapenet --data=${DATA} \ +# --resume=${PRETRAINED_MODEL} \ +# --gpus=${GPUS} --batch=${BATCH_SIZE} --gamma=0.3 + +# ##---------abo with 3D volume ----------- +# DATA=${BASE_DIR}/dataset_preprocessing/abo/abo_128_copy.zip +# # DATA=${BASE_DIR}/dataset_preprocessing/shapenet_cars/cars_128_copy.zip +# PRETRAINED_MODEL=${BASE_DIR}/eg3d/pretrained_models/shapenetcars128-64.pkl + +# python train.py --outdir=${BASE_DIR}/try-runs --cfg=abo_dataset --data=${DATA} \ +# --resume=${PRETRAINED_MODEL} \ +# --gpus=${GPUS} --batch=${BATCH_SIZE} --gamma=0.3 \ +# --backbone volume + + +##---------abo with 3D volume + no pretraining (because feature channel is down to 8)----------- +# DATA=${BASE_DIR}/dataset_preprocessing/shapenet_cars/cars_128_copy.zip +# DATA=${BASE_DIR}/dataset_preprocessing/abo/abo_128_copy.zip +# DATA=${BASE_DIR}/dataset_preprocessing/abo/abo_128_completed.zip +# DATA=${BASE_DIR}/dataset_preprocessing/abo/abo_128_completed_white.zip +# DATA=${BASE_DIR}/dataset_preprocessing/abo/abo_512_completed_white.zip +DATA=${BASE_DIR}/dataset_preprocessing/abo/abo_512_completed_white_small.zip +GPUS=2 +BATCH_SIZE=4 +python train.py --outdir=${BASE_DIR}/try-runs --cfg=abo_dataset --data=${DATA} \ + --gpus=${GPUS} --batch=${BATCH_SIZE} --gamma=0.3 \ + --backbone volume --decoder_dim 8 \ + --noise_strength 0.1 --snap 1 \ + --use_perception True --perception_reg 1 \ + --use_l2 True --l2_reg 1 \ + --use_chamfer True diff --git a/eg3d/torch_utils/.gitignore b/eg3d/torch_utils/.gitignore new file mode 100644 index 00000000..c0363794 --- /dev/null +++ b/eg3d/torch_utils/.gitignore @@ -0,0 +1 @@ +tmp/ \ No newline at end of file diff --git a/eg3d/torch_utils/chamfer3D/chamfer3D.cu b/eg3d/torch_utils/chamfer3D/chamfer3D.cu new file mode 100755 index 00000000..072dc086 --- /dev/null +++ b/eg3d/torch_utils/chamfer3D/chamfer3D.cu @@ -0,0 +1,196 @@ + +#include +#include + +#include +#include + +#include + + + +__global__ void NmDistanceKernel(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i){ + const int batch=512; + __shared__ float buf[batch*3]; + for (int i=blockIdx.x;ibest){ + result[(i*n+j)]=best; + result_i[(i*n+j)]=best_i; + } + } + __syncthreads(); + } + } +} +// int chamfer_cuda_forward(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i, cudaStream_t stream){ +int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2){ + + const auto batch_size = xyz1.size(0); + const auto n = dist1.size(1); //num_points point cloud A + const auto m = dist2.size(1); //num_points point cloud B + + NmDistanceKernel<<>>(batch_size, n, xyz1.data(), m, xyz2.data(), dist1.data(), idx1.data()); + NmDistanceKernel<<>>(batch_size, m, xyz2.data(), n, xyz1.data(), dist2.data(), idx2.data()); + + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + printf("error in nnd updateOutput: %s\n", cudaGetErrorString(err)); + //THError("aborting"); + return 0; + } + return 1; + + +} +__global__ void NmDistanceGradKernel(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,float * grad_xyz1,float * grad_xyz2){ + for (int i=blockIdx.x;i>>(batch_size,n,xyz1.data(),m,xyz2.data(),graddist1.data(),idx1.data(),gradxyz1.data(),gradxyz2.data()); + NmDistanceGradKernel<<>>(batch_size,m,xyz2.data(),n,xyz1.data(),graddist2.data(),idx2.data(),gradxyz2.data(),gradxyz1.data()); + + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + printf("error in nnd get grad: %s\n", cudaGetErrorString(err)); + //THError("aborting"); + return 0; + } + return 1; + +} + diff --git a/eg3d/torch_utils/chamfer3D/chamfer_cuda.cpp b/eg3d/torch_utils/chamfer3D/chamfer_cuda.cpp new file mode 100755 index 00000000..67574e21 --- /dev/null +++ b/eg3d/torch_utils/chamfer3D/chamfer_cuda.cpp @@ -0,0 +1,33 @@ +#include +#include + +///TMP +//#include "common.h" +/// NOT TMP + + +int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2); + + +int chamfer_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2); + + + + +int chamfer_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2) { + return chamfer_cuda_forward(xyz1, xyz2, dist1, dist2, idx1, idx2); +} + + +int chamfer_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, + at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2) { + + return chamfer_cuda_backward(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2); +} + + + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("forward", &chamfer_forward, "chamfer forward (CUDA)"); + m.def("backward", &chamfer_backward, "chamfer backward (CUDA)"); +} \ No newline at end of file diff --git a/eg3d/torch_utils/chamfer3D/dist_chamfer_3D.py b/eg3d/torch_utils/chamfer3D/dist_chamfer_3D.py new file mode 100644 index 00000000..de26d2c2 --- /dev/null +++ b/eg3d/torch_utils/chamfer3D/dist_chamfer_3D.py @@ -0,0 +1,81 @@ +from torch import nn +from torch.autograd import Function +import torch +import importlib +import os +chamfer_found = importlib.find_loader("chamfer_3D") is not None +if not chamfer_found: + ## Cool trick from https://github.com/chrdiller + print("Jitting Chamfer 3D") + cur_path = os.path.dirname(os.path.abspath(__file__)) + build_path = cur_path.replace('chamfer3D', 'tmp') + os.makedirs(build_path, exist_ok=True) + + from torch.utils.cpp_extension import load + chamfer_3D = load(name="chamfer_3D", + sources=[ + "/".join(os.path.abspath(__file__).split('/')[:-1] + ["chamfer_cuda.cpp"]), + "/".join(os.path.abspath(__file__).split('/')[:-1] + ["chamfer3D.cu"]), + ], build_directory=build_path) + print("Loaded JIT 3D CUDA chamfer distance") + +else: + import chamfer_3D + print("Loaded compiled 3D CUDA chamfer distance") + + +# Chamfer's distance module @thibaultgroueix +# GPU tensors only +class chamfer_3DFunction(Function): + @staticmethod + def forward(ctx, xyz1, xyz2): + batchsize, n, dim = xyz1.size() + assert dim==3, "Wrong last dimension for the chamfer distance 's input! Check with .size()" + _, m, dim = xyz2.size() + assert dim==3, "Wrong last dimension for the chamfer distance 's input! Check with .size()" + device = xyz1.device + + device = xyz1.device + + dist1 = torch.zeros(batchsize, n) + dist2 = torch.zeros(batchsize, m) + + idx1 = torch.zeros(batchsize, n).type(torch.IntTensor) + idx2 = torch.zeros(batchsize, m).type(torch.IntTensor) + + dist1 = dist1.to(device) + dist2 = dist2.to(device) + idx1 = idx1.to(device) + idx2 = idx2.to(device) + torch.cuda.set_device(device) + + chamfer_3D.forward(xyz1, xyz2, dist1, dist2, idx1, idx2) + ctx.save_for_backward(xyz1, xyz2, idx1, idx2) + return dist1, dist2, idx1, idx2 + + @staticmethod + def backward(ctx, graddist1, graddist2, gradidx1, gradidx2): + xyz1, xyz2, idx1, idx2 = ctx.saved_tensors + graddist1 = graddist1.contiguous() + graddist2 = graddist2.contiguous() + device = graddist1.device + + gradxyz1 = torch.zeros(xyz1.size()) + gradxyz2 = torch.zeros(xyz2.size()) + + gradxyz1 = gradxyz1.to(device) + gradxyz2 = gradxyz2.to(device) + chamfer_3D.backward( + xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2 + ) + return gradxyz1, gradxyz2 + + +class chamfer_3DDist(nn.Module): + def __init__(self): + super(chamfer_3DDist, self).__init__() + + def forward(self, input1, input2): + input1 = input1.contiguous() + input2 = input2.contiguous() + return chamfer_3DFunction.apply(input1, input2) diff --git a/eg3d/torch_utils/chamfer3D/setup.py b/eg3d/torch_utils/chamfer3D/setup.py new file mode 100755 index 00000000..9a23aada --- /dev/null +++ b/eg3d/torch_utils/chamfer3D/setup.py @@ -0,0 +1,14 @@ +from setuptools import setup +from torch.utils.cpp_extension import BuildExtension, CUDAExtension + +setup( + name='chamfer_3D', + ext_modules=[ + CUDAExtension('chamfer_3D', [ + "/".join(__file__.split('/')[:-1] + ['chamfer_cuda.cpp']), + "/".join(__file__.split('/')[:-1] + ['chamfer3D.cu']), + ]), + ], + cmdclass={ + 'build_ext': BuildExtension + }) \ No newline at end of file diff --git a/eg3d/torch_utils/misc.py b/eg3d/torch_utils/misc.py index 3c10e139..64c7ab45 100644 --- a/eg3d/torch_utils/misc.py +++ b/eg3d/torch_utils/misc.py @@ -15,6 +15,8 @@ import warnings import dnnlib +from ipdb import set_trace as st + #---------------------------------------------------------------------------- # Cached construction of constant tensors. Avoids CPU=>GPU copy when the # same constant is used multiple times. @@ -183,6 +185,8 @@ def check_ddp_consistency(module, ignore_regex=None): assert isinstance(module, torch.nn.Module) for name, tensor in named_params_and_buffers(module): fullname = type(module).__name__ + '.' + name + if 'backbone' in fullname: + continue if ignore_regex is not None and re.fullmatch(ignore_regex, fullname): continue tensor = tensor.detach() @@ -190,6 +194,7 @@ def check_ddp_consistency(module, ignore_regex=None): tensor = nan_to_num(tensor) other = tensor.clone() torch.distributed.broadcast(tensor=other, src=0) + assert (tensor == other).all(), fullname #---------------------------------------------------------------------------- diff --git a/eg3d/torch_utils/persistence.py b/eg3d/torch_utils/persistence.py index 1abf9cbf..c7136692 100644 --- a/eg3d/torch_utils/persistence.py +++ b/eg3d/torch_utils/persistence.py @@ -16,7 +16,8 @@ version of the code is not consistent with what was originally pickled.""" import sys -import pickle +# import pickle +import dill as pickle import io import inspect import copy diff --git a/eg3d/torch_utils/utils_ds.py b/eg3d/torch_utils/utils_ds.py new file mode 100644 index 00000000..005350ad --- /dev/null +++ b/eg3d/torch_utils/utils_ds.py @@ -0,0 +1,112 @@ +import numpy as np +import torch +import random +import torch.nn as nn + +from ipdb import set_trace as st + +def grp_range_torch(a,dev): + # st() + idx = torch.cumsum(a,0) + # st() + id_arr = torch.ones(idx[-1],dtype = torch.int64,device=dev) + id_arr[0] = 0 + # id_arr[idx[:-1]] = -a[:-1]+1 + try: + id_arr[idx[:-1]] = -a[:-1]+1 + except: + id_arr[idx[:]] = -a[:]+1 + return torch.cumsum(id_arr,0) + # generate array like [0,1,2,3,4,5,0,1,2,3,4,5,6] where each 0-n gives id to points inside the same grid + +def parallel_FPS(np_cat_fea,K): + return nb_greedy_FPS(np_cat_fea,K) + +def nb_greedy_FPS(xyz,K): + start_element = 0 + sample_num = xyz.shape[0] + sum_vec = np.zeros((sample_num,1),dtype = np.float32) + xyz_sq = xyz**2 + for j in range(sample_num): + sum_vec[j,0] = np.sum(xyz_sq[j,:]) + pairwise_distance = sum_vec + np.transpose(sum_vec) - 2*np.dot(xyz, np.transpose(xyz)) + + candidates_ind = np.zeros((sample_num,),dtype = np.bool_) + candidates_ind[start_element] = True + remain_ind = np.ones((sample_num,),dtype = np.bool_) + remain_ind[start_element] = False + all_ind = np.arange(sample_num) + + for i in range(1,K): + if i == 1: + min_remain_pt_dis = pairwise_distance[:,start_element] + min_remain_pt_dis = min_remain_pt_dis[remain_ind] + else: + cur_dis = pairwise_distance[remain_ind,:] + cur_dis = cur_dis[:,candidates_ind] + min_remain_pt_dis = np.zeros((cur_dis.shape[0],),dtype = np.float32) + for j in range(cur_dis.shape[0]): + min_remain_pt_dis[j] = np.min(cur_dis[j,:]) + next_ind_in_remain = np.argmax(min_remain_pt_dis) + next_ind = all_ind[remain_ind][next_ind_in_remain] + candidates_ind[next_ind] = True + remain_ind[next_ind] = False + + return candidates_ind + + +###########--------------------- +class Embedder: + def __init__(self, **kwargs): + self.kwargs = kwargs + self.create_embedding_fn() + + def create_embedding_fn(self): + embed_fns = [] + d = self.kwargs['input_dims'] + out_dim = 0 + if self.kwargs['include_input']: + embed_fns.append(lambda x : x) + out_dim += d + + max_freq = self.kwargs['max_freq_log2'] + N_freqs = self.kwargs['num_freqs'] + # st() + if self.kwargs['log_sampling']: + freq_bands = 2.**torch.linspace(0., max_freq, steps=N_freqs) + else: + freq_bands = torch.linspace(2.**0., 2.**max_freq, steps=N_freqs) + self.freq_bands = freq_bands.reshape(1,-1,1).cuda() + + for freq in freq_bands: + for p_fn in self.kwargs['periodic_fns']: + embed_fns.append(lambda x, p_fn=p_fn, freq=freq : p_fn(x * freq)) + out_dim += d + + self.embed_fns = embed_fns + self.out_dim = out_dim + + def embed(self, inputs): # inputs.shape [1024, 128, 3] + repeat = inputs.dim()-1 + inputs_scaled = (inputs.unsqueeze(-2) * self.freq_bands.contiguous().contiguous().view(*[1]*repeat,-1,1)).reshape(*inputs.shape[:-1],-1) + inputs_scaled = torch.cat((inputs, torch.sin(inputs_scaled), torch.cos(inputs_scaled)),dim=-1) + return inputs_scaled + +def get_embedder(multires, i=0, input_dims=3): + if i == -1: + return nn.Identity(), 3 + + embed_kwargs = { + 'include_input' : True, + 'input_dims' : input_dims, + 'max_freq_log2' : multires-1, + 'num_freqs' : multires, + 'log_sampling' : True, + 'periodic_fns' : [torch.sin, torch.cos], + } + + embedder_obj = Embedder(**embed_kwargs) + embed = lambda x, eo=embedder_obj : eo.embed(x) + return embed, embedder_obj.out_dim + # in case "pickle" failed, use below + return embedder_obj, embedder_obj.out_dim \ No newline at end of file diff --git a/eg3d/train.py b/eg3d/train.py index 7201e0ae..0d6c54f2 100644 --- a/eg3d/train.py +++ b/eg3d/train.py @@ -27,10 +27,13 @@ from torch_utils import training_stats from torch_utils import custom_ops +from ipdb import set_trace as st +import numpy as np + #---------------------------------------------------------------------------- def subprocess_fn(rank, c, temp_dir): - dnnlib.util.Logger(file_name=os.path.join(c.run_dir, 'log.txt'), file_mode='a', should_flush=True) + # dnnlib.util.Logger(file_name=os.path.join(c.run_dir, 'log.txt'), file_mode='a', should_flush=True) # Init torch.distributed. if c.num_gpus > 1: @@ -54,7 +57,7 @@ def subprocess_fn(rank, c, temp_dir): #---------------------------------------------------------------------------- def launch_training(c, desc, outdir, dry_run): - dnnlib.util.Logger(should_flush=True) + # dnnlib.util.Logger(should_flush=True) # Pick output directory. prev_run_dirs = [] @@ -193,6 +196,26 @@ def parse_comma_separated_list(s): @click.option('--reg_type', help='Type of regularization', metavar='STR', type=click.Choice(['l1', 'l1-alt', 'monotonic-detach', 'monotonic-fixed', 'total-variation']), required=False, default='l1') @click.option('--decoder_lr_mul', help='decoder learning rate multiplier.', metavar='FLOAT', type=click.FloatRange(min=0), default=1, required=False, show_default=True) +## specially for VolumeGenerator +@click.option('--backbone', help='whether use triplane or volume.', type=click.Choice(['triplane', 'volume']), required=False, default='triplane') +@click.option('--num_points', help='?.', metavar='INT', type=click.IntRange(min=512), required=False, default=1024) # default=1024 after finishing pipeline +@click.option('--num_materials', help='?.', metavar='INT', type=click.IntRange(min=3), required=False, default=9) +@click.option('--volume_res', help='volume resolution.', metavar='INT', type=click.IntRange(min=16), required=False, default=128) # default=128 after finishing pipeline +@click.option('--decoder_dim', help='OSGDecoder.', metavar='INT', type=click.IntRange(min=8), required=False, default=32) # default=128 after finishing pipeline +@click.option('--decoder_outdim', help='OSGDecoder.', metavar='INT', type=click.IntRange(min=8), required=False, default=32) # default=128 after finishing pipeline +@click.option('--use_ray_directions', help='If true, use_ray_directions during rendering.', metavar='BOOL', type=bool, required=False, default=True) +@click.option('--noise_strength', help='Control the magnitude of noises added to 3D volume during upsampling.', metavar='FLOAT', type=click.FloatRange(min=0, max=1), default=0.5, show_default=True) + +# specially for VolumeGenerator +# chamfer +@click.option('--use_chamfer', help='Use chamfer loss to regularize G', metavar='BOOL', type=bool, required=False, default=False) +@click.option('--chamfer_reg', help='chamfer reg', metavar='FLOAT', type=click.FloatRange(min=0.5), default=1, required=False, show_default=True) +@click.option('--use_perception', help='Use perception loss to regularize G', metavar='BOOL', type=bool, required=False, default=False) +@click.option('--perception_reg', help='perception reg', metavar='FLOAT', type=click.FloatRange(min=0.5), default=1, required=False, show_default=True) +@click.option('--use_l2', help='Use L2 loss to regularize G', metavar='BOOL', type=bool, required=False, default=False) +@click.option('--l2_reg', help='l2 reg', metavar='FLOAT', type=click.FloatRange(min=0.5), default=1, required=False, show_default=True) + + def main(**kwargs): """Train a GAN using the techniques described in the paper "Alias-Free Generative Adversarial Networks". @@ -218,6 +241,8 @@ def main(**kwargs): # Initialize config. opts = dnnlib.EasyDict(kwargs) # Command line arguments. + # st() + opts['mbstd_group']=1 # FIXME c = dnnlib.EasyDict() # Main config dict. c.G_kwargs = dnnlib.EasyDict(class_name=None, z_dim=512, w_dim=512, mapping_kwargs=dnnlib.EasyDict()) c.D_kwargs = dnnlib.EasyDict(class_name='training.networks_stylegan2.Discriminator', block_kwargs=dnnlib.EasyDict(), mapping_kwargs=dnnlib.EasyDict(), epilogue_kwargs=dnnlib.EasyDict()) @@ -243,6 +268,8 @@ def main(**kwargs): c.D_kwargs.block_kwargs.freeze_layers = opts.freezed c.D_kwargs.epilogue_kwargs.mbstd_group_size = opts.mbstd_group c.loss_kwargs.r1_gamma = opts.gamma + c.loss_kwargs.use_chamfer = opts.use_chamfer + c.loss_kwargs.chamfer_reg = opts.chamfer_reg c.G_opt_kwargs.lr = (0.002 if opts.cfg == 'stylegan2' else 0.0025) if opts.glr is None else opts.glr c.D_opt_kwargs.lr = opts.dlr c.metrics = opts.metrics @@ -264,7 +291,21 @@ def main(**kwargs): # Base configuration. c.ema_kimg = c.batch_size * 10 / 32 - c.G_kwargs.class_name = 'training.triplane.TriPlaneGenerator' + + ## conditional generator with pointcloud input + if opts.backbone == 'volume': + c.G_kwargs.class_name = 'training.volume.VolumeGenerator' + # c.G_kwargs.pc_dim = np.array([opts.num_points, opts.num_materials]) # num_pc or the num_material?? + c.G_kwargs.pc_dim = [opts.num_points, opts.num_materials] + c.G_kwargs.volume_res = opts.volume_res + c.G_kwargs.decoder_dim = opts.decoder_dim + c.G_kwargs.noise_strength = opts.noise_strength + # c.D_kwargs.class_name = 'training.volume_discriminator.VolumeDualDiscriminator' + # c.D_kwargs.class_name = 'training.dual_discriminator.DualDiscriminator' + # c.G_kwargs.decoder_outdim = opts.decoder_outdim + else: + c.G_kwargs.class_name = 'training.triplane.TriPlaneGenerator' + c.D_kwargs.class_name = 'training.dual_discriminator.DualDiscriminator' c.G_kwargs.fused_modconv_default = 'inference_only' # Speed up training by using regular convolutions instead of grouped convolutions. c.loss_kwargs.filter_mode = 'antialiased' # Filter mode for raw images ['antialiased', 'none', float [0-1]] @@ -296,6 +337,7 @@ def main(**kwargs): 'reg_type': opts.reg_type, # for experimenting with variations on density regularization 'decoder_lr_mul': opts.decoder_lr_mul, # learning rate multiplier for decoder 'sr_antialias': True, + 'use_ray_directions': opts.use_ray_directions, } if opts.cfg == 'ffhq': @@ -329,6 +371,19 @@ def main(**kwargs): 'avg_camera_radius': 1.7, 'avg_camera_pivot': [0, 0, 0], }) + elif opts.cfg == 'abo_dataset': + rendering_options.update({ + 'depth_resolution': 64, + 'depth_resolution_importance': 16, + 'ray_start': 0.1, + 'ray_end': 2.6, + 'box_warp': 1.6, + 'white_back': True, + 'avg_camera_radius': 1.7, + 'avg_camera_pivot': [0, 0, 0], + # 'decoder_output_dim': 8, + }) + else: assert False, "Need to specify config" @@ -352,6 +407,11 @@ def main(**kwargs): c.G_kwargs.sr_kwargs = dnnlib.EasyDict(channel_base=opts.cbase, channel_max=opts.cmax, fused_modconv_default='inference_only') c.loss_kwargs.style_mixing_prob = opts.style_mixing_prob + c.loss_kwargs.use_perception = opts.use_perception + c.loss_kwargs.perception_reg = opts.perception_reg + c.loss_kwargs.use_l2 = opts.use_l2 + c.loss_kwargs.l2_reg = opts.l2_reg + # Augmentation. if opts.aug != 'noaug': diff --git a/eg3d/training/chamfer_loss.py b/eg3d/training/chamfer_loss.py new file mode 100644 index 00000000..26bf187a --- /dev/null +++ b/eg3d/training/chamfer_loss.py @@ -0,0 +1,55 @@ +"""Chamfer Loss.""" + +import numpy as np +import torch +from torch_utils import persistence +from training.volumetric_rendering.ray_sampler import RaySampler + +from torch_utils.chamfer3D.dist_chamfer_3D import chamfer_3DDist +#---------------------------------------------------------------------------- +@persistence.persistent_class +class ChamferLoss(torch.nn.Module): + def __init__(self): + super().__init__() + self.ray_sampler = RaySampler() + self.chamfer3d = chamfer_3DDist() + + def forward(self, c, img, pc, neural_rendering_resolution): + dtype = torch.float32 + memory_format = torch.contiguous_format + B = c.shape[0] + loss_shape = (B ,1) + _device = img['image'].device + if pc is None or 'image_depth' not in img: + return torch.zeros(loss_shape).to(device=_device, dtype=dtype, memory_format=memory_format) + pc = pc[...,:3] + image_depth = img['image_depth'].view(img['image'].shape[0], -1, 1) + cam2world_matrix = c[:, :16].view(-1, 4, 4) + intrinsics = c[:, 16:25].view(-1, 3, 3) + + if neural_rendering_resolution is None: + neural_rendering_resolution = self.neural_rendering_resolution + else: + self.neural_rendering_resolution = neural_rendering_resolution + + # Create a batch of rays for volume rendering + ray_origins, ray_directions = self.ray_sampler(cam2world_matrix, intrinsics, neural_rendering_resolution) + # distance = ray_origins[:,0].unsqueeze(1).repeat(1, pc.shape[1], 1) + # distance = torch.sqrt(torch.sum((distance - pc) **2, axis=2)) + # max_distance, _ = torch.max(distance, axis = 1) + pred_pos = image_depth * ray_directions + ray_origins + # mask = image_depth.view(B, -1) < max_distance.view(B, -1) + chamfer_loss_0, chamfer_loss_1, _, _ = self.chamfer3d(pred_pos, pc) + chamfer_loss_0_sorted, _ = chamfer_loss_0.sort(1) + chamfer_loss_1_sorted, _ = chamfer_loss_1.sort(1) + chamfer_loss_0_sorted = chamfer_loss_0_sorted[:,:min(pred_pos.shape[1], pc.shape[1]) // 2] + chamfer_loss_1_sorted = chamfer_loss_0_sorted[:,:min(pred_pos.shape[1], pc.shape[1]) // 2] + chamfer_loss_0= torch.mean(chamfer_loss_0_sorted, dim=1).to(device=_device, dtype=dtype, memory_format=memory_format) + chamfer_loss_1= torch.mean(chamfer_loss_1_sorted, dim=1).to(device=_device, dtype=dtype, memory_format=memory_format) + chamfer_loss = chamfer_loss_1 + chamfer_loss_0 + # for pred_, mask_, pc_ in zip(pred_pos.split(1), mask.split(1), pc.split(1)): + # pred_ = pred_[mask_][None,...] + # _batch_chamfer_loss = self.chamfer3d(pred_, pc_)[self.direction] + # chamfer_loss += [torch.mean(_batch_chamfer_loss, dim=1).to(device=_device, dtype=dtype, memory_format=memory_format)] + + return chamfer_loss.view(loss_shape) diff --git a/eg3d/training/costregnet.py b/eg3d/training/costregnet.py new file mode 100644 index 00000000..8f652230 --- /dev/null +++ b/eg3d/training/costregnet.py @@ -0,0 +1,437 @@ +import torch +torch.autograd.set_detect_anomaly(True) +import torch.nn as nn +import torch.nn.functional as F +from inplace_abn import InPlaceABN + +from ipdb import set_trace as st + +#---------------------------------------- +class ConvBnReLU3D(nn.Module): + def __init__(self, in_channels, out_channels, + kernel_size=3, stride=1, pad=1, + norm_act=InPlaceABN): + super(ConvBnReLU3D, self).__init__() + self.conv = nn.Conv3d(in_channels, out_channels, + kernel_size, stride=stride, padding=pad, bias=False) + self.bn = norm_act(out_channels) + self.act = nn.LeakyReLU() + # self.bn = nn.ReLU() + # self.conv.apply(conv3d_weights_init) + + def forward(self, x): + return self.act(self.bn(self.conv(x))) + +class CostRegNet_Deeper(nn.Module): # 256^3 -> 8^3; 128^3 -> 4^3 + def __init__(self, in_channels, out_dim=8, norm_act=InPlaceABN): + super(CostRegNet_Deeper, self).__init__() + + self.conv0 = ConvBnReLU3D(in_channels, out_dim, norm_act=norm_act) + + self.conv1 = ConvBnReLU3D(out_dim, 16, stride=2, norm_act=norm_act) + self.conv2 = ConvBnReLU3D(16, 16, norm_act=norm_act) + + self.conv3 = ConvBnReLU3D(16, 32, stride=2, norm_act=norm_act) + self.conv4 = ConvBnReLU3D(32, 32, norm_act=norm_act) + + self.conv5 = ConvBnReLU3D(32, 64, stride=2, norm_act=norm_act) + self.conv6 = ConvBnReLU3D(64, 64, norm_act=norm_act) + + self.conv51 = ConvBnReLU3D(64, 64, stride=2, norm_act=norm_act) + self.conv61 = ConvBnReLU3D(64, 64, norm_act=norm_act) + + self.conv52 = ConvBnReLU3D(64, 64, stride=2, norm_act=norm_act) + self.conv62 = ConvBnReLU3D(64, 64, norm_act=norm_act) + + self.conv27 = nn.Sequential( + nn.ConvTranspose3d(64, 64, 3, padding=1, output_padding=1, + stride=2, bias=False), + norm_act(64)) + + self.conv17 = nn.Sequential( + nn.ConvTranspose3d(64, 64, 3, padding=1, output_padding=1, + stride=2, bias=False), + norm_act(64)) + + self.conv7 = nn.Sequential( + nn.ConvTranspose3d(64, 32, 3, padding=1, output_padding=1, + stride=2, bias=False), + norm_act(32)) + + self.conv9 = nn.Sequential( + nn.ConvTranspose3d(32, 16, 3, padding=1, output_padding=1, + stride=2, bias=False), + norm_act(16)) + + # self.conv11 = nn.Sequential( + # nn.ConvTranspose3d(16, 8, 3, padding=1, output_padding=1, + # stride=2, bias=False), + # norm_act(8)) + self.conv11 = nn.Sequential( + nn.ConvTranspose3d(16, out_dim, 3, padding=1, output_padding=1, + stride=2, bias=False), + norm_act(out_dim)) + + # self.conv12 = nn.Conv3d(8, 8, 3, stride=1, padding=1, bias=True) + + def forward(self, x): + conv0 = self.conv0(x) + + conv2 = self.conv2(self.conv1(conv0)) + conv4 = self.conv4(self.conv3(conv2)) + # if self.conv3.bn.weight.grad != None: + # st() + + # x = self.conv6(self.conv5(conv4)) + conv6 = self.conv6(self.conv5(conv4)) + + conv61 = self.conv61(self.conv51(conv6)) + conv62 = self.conv62(self.conv52(conv61)) + # print("CostRegNetDeeper bottleneck:", conv62.shape): # 256^3 -> 8^3; 128^3 -> 4^3 + x = conv61 + self.conv27(conv62) + x = conv6 + self.conv17(x) + + x = conv4 + self.conv7(x) + # del conv4 + x = conv2 + self.conv9(x) + # x = conv2 + self.conv9(conv4) + del conv2, conv4 + x = conv0 + self.conv11(x) + del conv0 + # x = self.conv12(x) + return x + + +class PcWsUnet(nn.Module): # 256^3 -> 8^3; 128^3 -> 4^3 + def __init__(self, in_channels, in_resolution, block_resolutions, out_dim=8, norm_act=InPlaceABN): + super(PcWsUnet, self).__init__() + self.block_resolutions = block_resolutions + self.conv0 = ConvBnReLU3D(in_channels, out_dim, norm_act=norm_act) + + self.conv1 = ConvBnReLU3D(out_dim, 16, stride=2, norm_act=norm_act) + self.conv2 = ConvBnReLU3D(16, 16, norm_act=norm_act) + + self.conv3 = ConvBnReLU3D(16, 32, stride=2, norm_act=norm_act) + self.conv4 = ConvBnReLU3D(32, 32, norm_act=norm_act) + + self.conv5 = ConvBnReLU3D(32, 64, stride=2, norm_act=norm_act) + self.conv6 = ConvBnReLU3D(64, 64, norm_act=norm_act) + + self.conv51 = ConvBnReLU3D(64, 64, stride=2, norm_act=norm_act) + self.conv61 = ConvBnReLU3D(64, 64, norm_act=norm_act) + + self.conv52 = ConvBnReLU3D(64, 64, stride=2, norm_act=norm_act) + self.conv62 = ConvBnReLU3D(64, 64, norm_act=norm_act) + + self.conv27 = nn.Sequential( + nn.ConvTranspose3d(64, 64, 3, padding=1, output_padding=1, + stride=2, bias=False), + norm_act(64)) + + self.conv17 = nn.Sequential( + nn.ConvTranspose3d(64, 64, 3, padding=1, output_padding=1, + stride=2, bias=False), + norm_act(64)) + + self.conv7 = nn.Sequential( + nn.ConvTranspose3d(64, 32, 3, padding=1, output_padding=1, + stride=2, bias=False), + norm_act(32)) + + self.conv9 = nn.Sequential( + nn.ConvTranspose3d(32, 16, 3, padding=1, output_padding=1, + stride=2, bias=False), + norm_act(16)) + + # self.conv11 = nn.Sequential( + # nn.ConvTranspose3d(16, 8, 3, padding=1, output_padding=1, + # stride=2, bias=False), + # norm_act(8)) + self.conv11 = nn.Sequential( + nn.ConvTranspose3d(16, out_dim, 3, padding=1, output_padding=1, + stride=2, bias=False), + norm_act(out_dim)) + + # self.conv12 = nn.Conv3d(8, 8, 3, stride=1, padding=1, bias=True) + + ## construct FC layers + max_res, min_res = max(self.block_resolutions), min(self.block_resolutions) + ## in_res --> outdim + ## inres//2 --> 16 + self.max_res = min(max_res, in_resolution, 32) + res = self.max_res + channels = {128:8, 64:16, 32:32, 16:64, 8:64, 4:64} + while (res>= min_res): + ch = channels.get(res) + layer = nn.Linear((res**3)*ch, ch*3) + setattr(self, f'fc{res}', layer) + # print(res, ch) + res = res//2 + + + def forward(self, x): + res_feature = {} + + conv0 = self.conv0(x) + # res_feature[conv0.shape[-1]]=conv0 + + conv2 = self.conv2(self.conv1(conv0)) + # res_feature[conv2.shape[-1]]=conv2 + + conv4 = self.conv4(self.conv3(conv2)) + # res_feature[conv4.shape[-1]]=conv4 + # if self.conv3.bn.weight.grad != None: + + # x = self.conv6(self.conv5(conv4)) + conv6 = self.conv6(self.conv5(conv4)) + # res_feature[conv6.shape[-1]]=conv6 + + conv61 = self.conv61(self.conv51(conv6)) + # res_feature[conv61.shape[-1]]=conv61 + + conv62 = self.conv62(self.conv52(conv61)) + # res_feature[conv62.shape[-1]]=conv62 + + # print("CostRegNetDeeper bottleneck:", conv62.shape): # 256^3 -> 8^3; 128^3 -> 4^3 + x = conv61 + self.conv27(conv62) + try: + res = x.shape[-1] + layer = getattr(self, f'fc{res}') + res_feature[x.shape[-1]]=layer(x.flatten(1)) + except: + pass + + x = conv6 + self.conv17(x) + try: + res = x.shape[-1] + layer = getattr(self, f'fc{res}') + res_feature[x.shape[-1]]=layer(x.flatten(1)) + except: + pass + # res_feature[x.shape[-1]]=x + + x = conv4 + self.conv7(x) + # res_feature[x.shape[-1]]=x + try: + res = x.shape[-1] + layer = getattr(self, f'fc{res}') + res_feature[x.shape[-1]]=layer(x.flatten(1)) + except: + pass + + # del conv4 + x = conv2 + self.conv9(x) + # res_feature[x.shape[-1]]=x + try: + res = x.shape[-1] + layer = getattr(self, f'fc{res}') + res_feature[x.shape[-1]]=layer(x.flatten(1)) + except: + pass + + # x = conv2 + self.conv9(conv4) + del conv2, conv4 + x = conv0 + self.conv11(x) + # if x.shape[-1] <= self.max_res: + # res_feature[x.shape[-1]]=x + try: + res = x.shape[-1] + layer = getattr(self, f'fc{res}') + res_feature[x.shape[-1]]=layer(x.flatten(1)) + except: + pass + + del conv0 + # x = self.conv12(x) + + ## use FC layer to project 3D volumes to 1d ws feature + # res_ws={} + # for res, vol in res_feature.items(): + # layer = getattr(self, f'fc{res}') + # ws = layer(vol.flatten(1)) + # res_feature.update({res:ws}) + + return res_feature + + +class Synthesis3DUnet(nn.Module): # 256^3 -> 8^3; 128^3 -> 4^3 + def __init__(self, + in_channels, + out_dim=8, + use_noise=False, + noise_strength = 0.5, + ws_channel=512, + affine_act='relu', #### ???? FIXME: is this a good activation + norm_act=InPlaceABN): + + super(Synthesis3DUnet, self).__init__() + + self.use_noise = use_noise + # noise_strength = 0.5 + self.noise_strength = noise_strength + + self.conv0 = ConvBnReLU3D(in_channels, out_dim, norm_act=norm_act) + + self.conv1 = ConvBnReLU3D(out_dim, 16, stride=2, norm_act=norm_act) + self.conv2 = ConvBnReLU3D(16, 16, norm_act=norm_act) + + self.conv3 = ConvBnReLU3D(16, 32, stride=2, norm_act=norm_act) + self.conv4 = ConvBnReLU3D(32, 32, norm_act=norm_act) + + self.conv5 = ConvBnReLU3D(32, 64, stride=2, norm_act=norm_act) + self.conv6 = ConvBnReLU3D(64, 64, norm_act=norm_act) + + self.conv51 = ConvBnReLU3D(64, 64, stride=2, norm_act=norm_act) + self.conv61 = ConvBnReLU3D(64, 64, norm_act=norm_act) + + self.conv52 = ConvBnReLU3D(64, 64, stride=2, norm_act=norm_act) + self.conv62 = ConvBnReLU3D(64, 64, norm_act=norm_act) + + self.conv27 = nn.Sequential( + nn.ConvTranspose3d(64, 64, 3, padding=1, output_padding=1, + stride=2, bias=False), + norm_act(64)) + self.affine27 = nn.Sequential( + nn.Linear(ws_channel, 64), + nn.ReLU() + ) + + + self.conv17 = nn.Sequential( + nn.ConvTranspose3d(64, 64, 3, padding=1, output_padding=1, + stride=2, bias=False), + norm_act(64)) + self.affine17 = nn.Sequential( + nn.Linear(ws_channel, 64), + nn.ReLU() + ) + + self.conv7 = nn.Sequential( + nn.ConvTranspose3d(64, 32, 3, padding=1, output_padding=1, + stride=2, bias=False), + norm_act(32)) + self.affine7 = nn.Sequential( + nn.Linear(ws_channel, 32), + nn.ReLU() + ) + + self.conv9 = nn.Sequential( + nn.ConvTranspose3d(32, 16, 3, padding=1, output_padding=1, + stride=2, bias=False), + norm_act(16)) + self.affine9 = nn.Sequential( + nn.Linear(ws_channel, 16), + nn.ReLU() + ) + + # self.conv11 = nn.Sequential( + # nn.ConvTranspose3d(16, 8, 3, padding=1, output_padding=1, + # stride=2, bias=False), + # norm_act(8)) + self.conv11 = nn.Sequential( + nn.ConvTranspose3d(16, out_dim, 3, padding=1, output_padding=1, + stride=2, bias=False), + norm_act(out_dim)) + self.affine11 = nn.Sequential( + nn.Linear(ws_channel, out_dim), + nn.ReLU() + ) + + + # self.conv12 = nn.Conv3d(8, 8, 3, stride=1, padding=1, bias=True) + + def forward(self, x, ws): + conv0 = self.conv0(x) + + + conv2 = self.conv2(self.conv1(conv0)) + conv4 = self.conv4(self.conv3(conv2)) + # if self.conv3.bn.weight.grad != None: + # st() + + # x = self.conv6(self.conv5(conv4)) + conv6 = self.conv6(self.conv5(conv4)) + + conv61 = self.conv61(self.conv51(conv6)) + conv62 = self.conv62(self.conv52(conv61)) + # print("CostRegNetDeeper bottleneck:", conv62.shape) # 256^3 -> 8^3; 128^3 -> 4^3 + + ### below is upconv process: add noises and latent + w_idx=0 + + + x = conv61 + self.conv27(conv62) + if self.use_noise: + noise = torch.rand(x.shape, device=x.device, dtype=torch.float32) + noise = noise*self.noise_strength + + x = x.add_(noise.to(x.dtype)) + style = self.affine27(ws.narrow(1, w_idx, 1)).permute(0,2,1) + w_idx += 1 + # st() + assert x.shape[:2] == style.shape[:2] + B, C = x.shape[:2] + style = style.reshape(B,C,1,1,1) # extend to 3D + x = x*style + # st() # x.shape + + + x = conv6 + self.conv17(x) + if self.use_noise: + noise = torch.rand(x.shape, device=x.device) + noise = noise*self.noise_strength + x = x.add_(noise.to(x.dtype)) + style = self.affine17(ws.narrow(1, w_idx, 1)).permute(0,2,1) + w_idx += 1 + assert x.shape[:2] == style.shape[:2] + B, C = x.shape[:2] + style = style.reshape(B,C,1,1,1) # extend to 3D + x = x*style + + + x = conv4 + self.conv7(x) + if self.use_noise: + noise = torch.rand(x.shape, device=x.device) + noise = noise*self.noise_strength + x = x.add_(noise.to(x.dtype)) + style = self.affine7(ws.narrow(1, w_idx, 1)).permute(0,2,1) + w_idx += 1 + assert x.shape[:2] == style.shape[:2] + B, C = x.shape[:2] + style = style.reshape(B,C,1,1,1) # extend to 3D + x = x*style + # del conv4 + + + x = conv2 + self.conv9(x) + if self.use_noise: + noise = torch.rand(x.shape, device=x.device) + noise = noise*self.noise_strength + x = x.add_(noise.to(x.dtype)) + style = self.affine9(ws.narrow(1, w_idx, 1)).permute(0,2,1) + w_idx += 1 + assert x.shape[:2] == style.shape[:2] + B, C = x.shape[:2] + style = style.reshape(B,C,1,1,1) # extend to 3D + x = x*style + + del conv2, conv4 + + + x = conv0 + self.conv11(x) + if self.use_noise: + noise = torch.rand(x.shape, device=x.device) + noise = noise*self.noise_strength + x = x.add_(noise.to(x.dtype)) + style = self.affine11(ws.narrow(1, w_idx, 1)).permute(0,2,1) + w_idx += 1 + assert x.shape[:2] == style.shape[:2] + B, C = x.shape[:2] + style = style.reshape(B,C,1,1,1) # extend to 3D + x = x*style + + del conv0 + + # print(f"Totally used up to {w_idx} ws in synthesis3DUnet") ## currently used only 5 + + return x \ No newline at end of file diff --git a/eg3d/training/dataset.py b/eg3d/training/dataset.py index b4d7c4fb..c5cd9651 100644 --- a/eg3d/training/dataset.py +++ b/eg3d/training/dataset.py @@ -17,12 +17,14 @@ import json import torch import dnnlib +from ipdb import set_trace as st try: import pyspng except ImportError: pyspng = None +import pandas as pd #---------------------------------------------------------------------------- class Dataset(torch.utils.data.Dataset): @@ -54,7 +56,8 @@ def __init__(self, def _get_raw_labels(self): if self._raw_labels is None: - self._raw_labels = self._load_raw_labels() if self._use_labels else None + # self._raw_labels = self._load_raw_labels() if self._use_labels else None + self._raw_labels = self._load_raw_labels() if self._raw_labels is None: self._raw_labels = np.zeros([self._raw_shape[0], 0], dtype=np.float32) assert isinstance(self._raw_labels, np.ndarray) @@ -73,7 +76,9 @@ def _load_raw_image(self, raw_idx): # to be overridden by subclass raise NotImplementedError def _load_raw_labels(self): # to be overridden by subclass - raise NotImplementedError + st() + print('to be overridden by subclass') + # raise NotImplementedError def __getstate__(self): return dict(self.__dict__, _raw_labels=None) @@ -95,7 +100,12 @@ def __getitem__(self, idx): if self._xflip[idx]: assert image.ndim == 3 # CHW image = image[:, :, ::-1] - return image.copy(), self.get_label(idx) + try: + pointcloud = self._load_raw_pointcloud(self._raw_idx[idx]) + except: + pointcloud = np.empty([1,0]) + + return image.copy(), self.get_label(idx), pointcloud.copy() def get_label(self, idx): label = self._get_raw_labels()[self._raw_idx[idx]] @@ -104,6 +114,10 @@ def get_label(self, idx): onehot[label] = 1 label = onehot return label.copy() + + def get_pointcloud(self, idx): + pointcloud = self._load_raw_pointcloud(self._raw_idx[idx]) + return pointcloud.copy() def get_details(self, idx): d = dnnlib.EasyDict() @@ -136,6 +150,7 @@ def resolution(self): @property def label_shape(self): + # st() if self._label_shape is None: raw_labels = self._get_raw_labels() if raw_labels.dtype == np.int64: @@ -151,6 +166,7 @@ def label_dim(self): @property def has_labels(self): + return any(x != 0 for x in self.label_shape) @property @@ -173,12 +189,15 @@ def __init__(self, self._all_fnames = {os.path.relpath(os.path.join(root, fname), start=self._path) for root, _dirs, files in os.walk(self._path) for fname in files} elif self._file_ext(self._path) == '.zip': self._type = 'zip' + # st() self._all_fnames = set(self._get_zipfile().namelist()) else: raise IOError('Path must point to a directory or zip') PIL.Image.init() self._image_fnames = sorted(fname for fname in self._all_fnames if self._file_ext(fname) in PIL.Image.EXTENSION) + self._pc_fnames = sorted(fname for fname in self._all_fnames if self._file_ext(fname) == '.csv') + if len(self._image_fnames) == 0: raise IOError('No image files found in the specified path') @@ -187,6 +206,7 @@ def __init__(self, if resolution is not None and (raw_shape[2] != resolution or raw_shape[3] != resolution): raise IOError('Image files do not match the specified resolution') super().__init__(name=name, raw_shape=raw_shape, **super_kwargs) + # st() @staticmethod def _file_ext(fname): @@ -226,8 +246,26 @@ def _load_raw_image(self, raw_idx): image = image[:, :, np.newaxis] # HW => HWC image = image.transpose(2, 0, 1) # HWC => CHW return image + + def _load_raw_pointcloud(self, raw_idx): + fname = self._pc_fnames[raw_idx] + + with self._open_file(fname) as f: + df = pd.read_csv(f, header=None) + pc_array = df.values.astype(np.float32) + return pc_array + + def get_label_and_pc(self, idx): + label = self._get_raw_labels()[self._raw_idx[idx]] + if label.dtype == np.int64: + onehot = np.zeros(self.label_shape, dtype=np.float32) + onehot[label] = 1 + label = onehot + pointcloud = self._load_raw_pointcloud(self._raw_idx[idx]) + return label.copy(), pointcloud.copy() def _load_raw_labels(self): + fname = 'dataset.json' if fname not in self._all_fnames: return None @@ -235,6 +273,7 @@ def _load_raw_labels(self): labels = json.load(f)['labels'] if labels is None: return None + labels = dict(labels) labels = [labels[fname.replace('\\', '/')] for fname in self._image_fnames] labels = np.array(labels) diff --git a/eg3d/training/loss.py b/eg3d/training/loss.py index b2c637a6..3a2a0126 100644 --- a/eg3d/training/loss.py +++ b/eg3d/training/loss.py @@ -16,7 +16,14 @@ from torch_utils.ops import conv2d_gradfix from torch_utils.ops import upfirdn2d from training.dual_discriminator import filtered_resizing +from training.chamfer_loss import ChamferLoss +from ipdb import set_trace as st +from training.volume import VolumeGenerator + +import clip +import torchvision.transforms as T +from PIL import Image #---------------------------------------------------------------------------- class Loss: @@ -26,7 +33,8 @@ def accumulate_gradients(self, phase, real_img, real_c, gen_z, gen_c, gain, cur_ #---------------------------------------------------------------------------- class StyleGAN2Loss(Loss): - def __init__(self, device, G, D, augment_pipe=None, r1_gamma=10, style_mixing_prob=0, pl_weight=0, pl_batch_shrink=2, pl_decay=0.01, pl_no_weight_grad=False, blur_init_sigma=0, blur_fade_kimg=0, r1_gamma_init=0, r1_gamma_fade_kimg=0, neural_rendering_resolution_initial=64, neural_rendering_resolution_final=None, neural_rendering_resolution_fade_kimg=0, gpc_reg_fade_kimg=1000, gpc_reg_prob=None, dual_discrimination=False, filter_mode='antialiased'): + def __init__(self, device, G, D, augment_pipe=None, r1_gamma=10, style_mixing_prob=0, pl_weight=0, pl_batch_shrink=2, pl_decay=0.01, pl_no_weight_grad=False, blur_init_sigma=0, blur_fade_kimg=0, r1_gamma_init=0, r1_gamma_fade_kimg=0, neural_rendering_resolution_initial=64, neural_rendering_resolution_final=None, neural_rendering_resolution_fade_kimg=0, gpc_reg_fade_kimg=1000, gpc_reg_prob=None, dual_discrimination=False, filter_mode='antialiased', + use_perception=False, perception_reg=1, use_l2=False, l2_reg=1, use_chamfer=False, chamfer_reg=1): super().__init__() self.device = device self.G = G @@ -52,9 +60,28 @@ def __init__(self, device, G, D, augment_pipe=None, r1_gamma=10, style_mixing_pr self.filter_mode = filter_mode self.resample_filter = upfirdn2d.setup_filter([1,3,3,1], device=device) self.blur_raw_target = True + assert self.gpc_reg_prob is None or (0 <= self.gpc_reg_prob <= 1) - - def run_G(self, z, c, swapping_prob, neural_rendering_resolution, update_emas=False): + self.use_perception = use_perception + self.perception_reg = perception_reg + if self.use_perception: + # device = "cuda" if torch.cuda.is_available() else "cpu" + # self.perception_reg = 1 # ViT will give larger loss + # model, preprocess = clip.load("RN50", device=device) + self.clip_model, self.clip_preprocess = clip.load("ViT-B/32", device=device) # maybe too large + for p in self.clip_model.parameters(): + p.requires_grad=False + + self.use_l2 = use_l2 + self.l2_reg = l2_reg + + self.use_chamfer = use_chamfer + if self.use_chamfer: + self.chamfer_reg = chamfer_reg + self.chamfer_loss = ChamferLoss() + + + def run_G(self, z, c, pc, swapping_prob, neural_rendering_resolution, update_emas=False): if swapping_prob is not None: c_swapped = torch.roll(c.clone(), 1, 0) c_gen_conditioning = torch.where(torch.rand((c.shape[0], 1), device=c.device) < swapping_prob, c_swapped, c) @@ -67,10 +94,14 @@ def run_G(self, z, c, swapping_prob, neural_rendering_resolution, update_emas=Fa cutoff = torch.empty([], dtype=torch.int64, device=ws.device).random_(1, ws.shape[1]) cutoff = torch.where(torch.rand([], device=ws.device) < self.style_mixing_prob, cutoff, torch.full_like(cutoff, ws.shape[1])) ws[:, cutoff:] = self.G.mapping(torch.randn_like(z), c, update_emas=False)[:, cutoff:] - gen_output = self.G.synthesis(ws, c, neural_rendering_resolution=neural_rendering_resolution, update_emas=update_emas) + # st() + if isinstance(self.G, VolumeGenerator): + gen_output = self.G.synthesis(ws, c, pc, neural_rendering_resolution=neural_rendering_resolution, update_emas=update_emas) + else: + gen_output = self.G.synthesis(ws, c, neural_rendering_resolution=neural_rendering_resolution, update_emas=update_emas) return gen_output, ws - def run_D(self, img, c, blur_sigma=0, blur_sigma_raw=0, update_emas=False): + def run_D(self, img, c, blur_sigma=0, blur_sigma_raw=0, update_emas=False): blur_size = np.floor(blur_sigma * 3) if blur_size > 0: with torch.autograd.profiler.record_function('blur'): @@ -86,16 +117,61 @@ def run_D(self, img, c, blur_sigma=0, blur_sigma_raw=0, update_emas=False): logits = self.D(img, c, update_emas=update_emas) return logits - - def accumulate_gradients(self, phase, real_img, real_c, gen_z, gen_c, gain, cur_nimg): + + def cal_perception_loss(self, gen_img, real_img): + + # convert to PIL + transform = T.ToPILImage() + + gen_img_PIL=[] + # currently all images are normalized within -1~1 + for gi in gen_img['image']: + # gi = (gi+1.)/2 # no difference + gen_img_PIL.append(self.clip_preprocess(transform(gi))) + gen_img_processed = torch.tensor(np.stack(gen_img_PIL)).to(self.device) + + real_img_PIL=[] + for ri in real_img['image']: + # ri = (ri+1.)/2 # no difference + real_img_PIL.append(self.clip_preprocess(transform(ri))) + real_img_processed = torch.tensor(np.stack(real_img_PIL)).to(self.device) + + # image = preprocess(Image.open("CLIP.png")).unsqueeze(0).to(device) + # text = clip.tokenize(["a diagram", "a dog", "a cat"]).to(device) + + gen_image_features = self.clip_model.encode_image(gen_img_processed).float() + real_image_features = self.clip_model.encode_image(real_img_processed).float() + + ## below MSUT NOT in no_grad! + mse = torch.nn.MSELoss(reduction='none') + loss = mse(gen_image_features, real_image_features) + + return loss + + + def cal_l2_loss(self, gen_img, real_img): + + gen_image_features = gen_img['image'] + real_image_features = real_img['image'] + + mse = torch.nn.MSELoss(reduction='none') + loss = mse(gen_image_features, real_image_features) + + return loss + + def accumulate_gradients(self, phase, real_img, real_c, gen_z, gen_c, gen_pc, gain, cur_nimg): assert phase in ['Gmain', 'Greg', 'Gboth', 'Dmain', 'Dreg', 'Dboth'] + ############# FIXME Oct 25: uncomment below to still enable Greg phase ################## if self.G.rendering_kwargs.get('density_reg', 0) == 0: phase = {'Greg': 'none', 'Gboth': 'Gmain'}.get(phase, phase) + ############################### if self.r1_gamma == 0: phase = {'Dreg': 'none', 'Dboth': 'Dmain'}.get(phase, phase) blur_sigma = max(1 - cur_nimg / (self.blur_fade_kimg * 1e3), 0) * self.blur_init_sigma if self.blur_fade_kimg > 0 else 0 r1_gamma = self.r1_gamma + + alpha = min(cur_nimg / (self.gpc_reg_fade_kimg * 1e3), 1) if self.gpc_reg_fade_kimg > 0 else 1 swapping_prob = (1 - alpha) * 1 + alpha * self.gpc_reg_prob if self.gpc_reg_prob is not None else None @@ -118,22 +194,68 @@ def accumulate_gradients(self, phase, real_img, real_c, gen_z, gen_c, gain, cur_ # Gmain: Maximize logits for generated images. if phase in ['Gmain', 'Gboth']: with torch.autograd.profiler.record_function('Gmain_forward'): - gen_img, _gen_ws = self.run_G(gen_z, gen_c, swapping_prob=swapping_prob, neural_rendering_resolution=neural_rendering_resolution) - gen_logits = self.run_D(gen_img, gen_c, blur_sigma=blur_sigma) + gen_img, _gen_ws = self.run_G(gen_z, gen_c, gen_pc, swapping_prob=swapping_prob, neural_rendering_resolution=neural_rendering_resolution) + + gen_logits = self.run_D(gen_img, gen_c, blur_sigma=blur_sigma) training_stats.report('Loss/scores/fake', gen_logits) training_stats.report('Loss/signs/fake', gen_logits.sign()) loss_Gmain = torch.nn.functional.softplus(-gen_logits) training_stats.report('Loss/G/loss', loss_Gmain) + + # chamfer loss + if self.use_chamfer: + chamfer_loss = self.chamfer_loss(gen_c, gen_img, gen_pc, neural_rendering_resolution) + chamfer_loss *= self.chamfer_reg + loss_Gmain += chamfer_loss + print(f"---------loss_chamfer\t(x{self.chamfer_reg}): {(chamfer_loss).sum().item()}-------------") + training_stats.report('Loss/G/chamfer_loss', chamfer_loss) + + # perceptual loss + if self.use_perception: + perception_loss = self.cal_perception_loss(gen_img=gen_img, real_img=real_img) + perception_loss = torch.mean(perception_loss, 1, True) * self.perception_reg + loss_Gmain += perception_loss + print(f"---------loss_perception\t(x{self.perception_reg}): {(perception_loss).sum().item()}-------------") + + training_stats.report('Loss/G/perceptual_loss', perception_loss) + + # L2 loss on the whole gen image + if self.use_l2: + l2_loss = self.cal_l2_loss(gen_img=gen_img, real_img=real_img) + l2_loss = torch.mean(l2_loss.flatten(1), -1, True) * self.l2_reg + loss_Gmain += l2_loss + print(f"---------loss_l2\t\t(x{self.l2_reg}): {(l2_loss).sum().item()}-------------") + + training_stats.report('Loss/G/l2_loss_whole', l2_loss) + + with torch.autograd.profiler.record_function('Gmain_backward'): loss_Gmain.mean().mul(gain).backward() + + + # Density Regularization if phase in ['Greg', 'Gboth'] and self.G.rendering_kwargs.get('density_reg', 0) > 0 and self.G.rendering_kwargs['reg_type'] == 'l1': - if swapping_prob is not None: - c_swapped = torch.roll(gen_c.clone(), 1, 0) - c_gen_conditioning = torch.where(torch.rand([], device=gen_c.device) < swapping_prob, c_swapped, gen_c) + + + if swapping_prob is not None: # always None + # c_swapped = torch.roll(gen_c.clone(), 1, 0) + # c_gen_conditioning = torch.where(torch.rand([], device=gen_c.device) < swapping_prob, c_swapped, gen_c) + ######## align pc_gen with c_gen ######### + B, N, pc_dim = gen_pc.shape + st() + gen_c_and_pc = torch.cat([gen_c, gen_pc.reshape(B,-1)], dim=-1) + c_and_pc_swapped = torch.roll(gen_c_and_pc.clone(), 1, 0) + c_and_pc_gen_conditioning = torch.where(torch.rand([], device=gen_c_and_pc.device) < swapping_prob, c_and_pc_swapped, gen_c_and_pc) + c_gen_conditioning = c_and_pc_gen_conditioning[...,:-pc_dim] + pc_gen_conditioning = c_and_pc_gen_conditioning[...,-pc_dim:].reshape(B,N,pc_dim) + assert (pc_gen_conditioning.shape==gen_pc.shape) and (c_gen_conditioning.shape==gen_c.shape) + st() + else: c_gen_conditioning = torch.zeros_like(gen_c) + pc_gen_conditioning = gen_pc ws = self.G.mapping(gen_z, c_gen_conditioning, update_emas=False) if self.style_mixing_prob > 0: @@ -144,7 +266,13 @@ def accumulate_gradients(self, phase, real_img, real_c, gen_z, gen_c, gain, cur_ initial_coordinates = torch.rand((ws.shape[0], 1000, 3), device=ws.device) * 2 - 1 perturbed_coordinates = initial_coordinates + torch.randn_like(initial_coordinates) * self.G.rendering_kwargs['density_reg_p_dist'] all_coordinates = torch.cat([initial_coordinates, perturbed_coordinates], dim=1) - sigma = self.G.sample_mixed(all_coordinates, torch.randn_like(all_coordinates), ws, update_emas=False)['sigma'] + #### FIXME: the coordinates fed into mixed_sample are both random, why is this????? + if isinstance(self.G, VolumeGenerator): + sigma = self.G.sample_mixed(all_coordinates, torch.randn_like(all_coordinates), ws, pc=pc_gen_conditioning,\ + box_warp=self.G.rendering_kwargs['box_warp'], update_emas=False)['sigma'] + # st() + else: + sigma = self.G.sample_mixed(all_coordinates, torch.randn_like(all_coordinates), ws, update_emas=False)['sigma'] sigma_initial = sigma[:, :sigma.shape[1]//2] sigma_perturbed = sigma[:, sigma.shape[1]//2:] @@ -158,6 +286,7 @@ def accumulate_gradients(self, phase, real_img, real_c, gen_z, gen_c, gain, cur_ c_gen_conditioning = torch.where(torch.rand([], device=gen_c.device) < swapping_prob, c_swapped, gen_c) else: c_gen_conditioning = torch.zeros_like(gen_c) + pc_gen_conditioning = gen_pc ws = self.G.mapping(gen_z, c_gen_conditioning, update_emas=False) @@ -165,7 +294,12 @@ def accumulate_gradients(self, phase, real_img, real_c, gen_z, gen_c, gain, cur_ perturbed_coordinates = initial_coordinates + torch.tensor([0, 0, -1], device=ws.device) * (1/256) * self.G.rendering_kwargs['box_warp'] # Behind all_coordinates = torch.cat([initial_coordinates, perturbed_coordinates], dim=1) - sigma = self.G.sample_mixed(all_coordinates, torch.randn_like(all_coordinates), ws, update_emas=False)['sigma'] + if isinstance(self.G, VolumeGenerator): + sigma = self.G.sample_mixed(all_coordinates, torch.randn_like(all_coordinates), ws, pc=pc_gen_conditioning,\ + box_warp=self.G.rendering_kwargs['box_warp'], update_emas=False)['sigma'] + # st() + else: + sigma = self.G.sample_mixed(all_coordinates, torch.randn_like(all_coordinates), ws, update_emas=False)['sigma'] sigma_initial = sigma[:, :sigma.shape[1]//2] sigma_perturbed = sigma[:, sigma.shape[1]//2:] @@ -178,6 +312,7 @@ def accumulate_gradients(self, phase, real_img, real_c, gen_z, gen_c, gain, cur_ c_gen_conditioning = torch.where(torch.rand([], device=gen_c.device) < swapping_prob, c_swapped, gen_c) else: c_gen_conditioning = torch.zeros_like(gen_c) + pc_gen_conditioning = gen_pc ws = self.G.mapping(gen_z, c_gen_conditioning, update_emas=False) if self.style_mixing_prob > 0: @@ -188,7 +323,12 @@ def accumulate_gradients(self, phase, real_img, real_c, gen_z, gen_c, gain, cur_ initial_coordinates = torch.rand((ws.shape[0], 1000, 3), device=ws.device) * 2 - 1 perturbed_coordinates = initial_coordinates + torch.randn_like(initial_coordinates) * (1/256) * self.G.rendering_kwargs['box_warp'] all_coordinates = torch.cat([initial_coordinates, perturbed_coordinates], dim=1) - sigma = self.G.sample_mixed(all_coordinates, torch.randn_like(all_coordinates), ws, update_emas=False)['sigma'] + if isinstance(self.G, VolumeGenerator): + sigma = self.G.sample_mixed(all_coordinates, torch.randn_like(all_coordinates), ws, pc=pc_gen_conditioning,\ + box_warp=self.G.rendering_kwargs['box_warp'], update_emas=False)['sigma'] + # st() + else: + sigma = self.G.sample_mixed(all_coordinates, torch.randn_like(all_coordinates), ws, update_emas=False)['sigma'] sigma_initial = sigma[:, :sigma.shape[1]//2] sigma_perturbed = sigma[:, sigma.shape[1]//2:] @@ -202,6 +342,7 @@ def accumulate_gradients(self, phase, real_img, real_c, gen_z, gen_c, gain, cur_ c_gen_conditioning = torch.where(torch.rand([], device=gen_c.device) < swapping_prob, c_swapped, gen_c) else: c_gen_conditioning = torch.zeros_like(gen_c) + pc_gen_conditioning = gen_pc ws = self.G.mapping(gen_z, c_gen_conditioning, update_emas=False) @@ -209,7 +350,12 @@ def accumulate_gradients(self, phase, real_img, real_c, gen_z, gen_c, gain, cur_ perturbed_coordinates = initial_coordinates + torch.tensor([0, 0, -1], device=ws.device) * (1/256) * self.G.rendering_kwargs['box_warp'] # Behind all_coordinates = torch.cat([initial_coordinates, perturbed_coordinates], dim=1) - sigma = self.G.sample_mixed(all_coordinates, torch.randn_like(all_coordinates), ws, update_emas=False)['sigma'] + if isinstance(self.G, VolumeGenerator): + sigma = self.G.sample_mixed(all_coordinates, torch.randn_like(all_coordinates), ws, pc=pc_gen_conditioning,\ + box_warp=self.G.rendering_kwargs['box_warp'], update_emas=False)['sigma'] + # st() + else: + sigma = self.G.sample_mixed(all_coordinates, torch.randn_like(all_coordinates), ws, update_emas=False)['sigma'] sigma_initial = sigma[:, :sigma.shape[1]//2] sigma_perturbed = sigma[:, sigma.shape[1]//2:] @@ -222,6 +368,7 @@ def accumulate_gradients(self, phase, real_img, real_c, gen_z, gen_c, gain, cur_ c_gen_conditioning = torch.where(torch.rand([], device=gen_c.device) < swapping_prob, c_swapped, gen_c) else: c_gen_conditioning = torch.zeros_like(gen_c) + pc_gen_conditioning = gen_pc ws = self.G.mapping(gen_z, c_gen_conditioning, update_emas=False) if self.style_mixing_prob > 0: @@ -232,7 +379,12 @@ def accumulate_gradients(self, phase, real_img, real_c, gen_z, gen_c, gain, cur_ initial_coordinates = torch.rand((ws.shape[0], 1000, 3), device=ws.device) * 2 - 1 perturbed_coordinates = initial_coordinates + torch.randn_like(initial_coordinates) * (1/256) * self.G.rendering_kwargs['box_warp'] all_coordinates = torch.cat([initial_coordinates, perturbed_coordinates], dim=1) - sigma = self.G.sample_mixed(all_coordinates, torch.randn_like(all_coordinates), ws, update_emas=False)['sigma'] + if isinstance(self.G, VolumeGenerator): + sigma = self.G.sample_mixed(all_coordinates, torch.randn_like(all_coordinates), ws, pc=pc_gen_conditioning,\ + box_warp=self.G.rendering_kwargs['box_warp'], update_emas=False)['sigma'] + # st() + else: + sigma = self.G.sample_mixed(all_coordinates, torch.randn_like(all_coordinates), ws, update_emas=False)['sigma'] sigma_initial = sigma[:, :sigma.shape[1]//2] sigma_perturbed = sigma[:, sigma.shape[1]//2:] @@ -243,7 +395,7 @@ def accumulate_gradients(self, phase, real_img, real_c, gen_z, gen_c, gain, cur_ loss_Dgen = 0 if phase in ['Dmain', 'Dboth']: with torch.autograd.profiler.record_function('Dgen_forward'): - gen_img, _gen_ws = self.run_G(gen_z, gen_c, swapping_prob=swapping_prob, neural_rendering_resolution=neural_rendering_resolution, update_emas=True) + gen_img, _gen_ws = self.run_G(gen_z, gen_c, gen_pc, swapping_prob=swapping_prob, neural_rendering_resolution=neural_rendering_resolution, update_emas=True) gen_logits = self.run_D(gen_img, gen_c, blur_sigma=blur_sigma, update_emas=True) training_stats.report('Loss/scores/fake', gen_logits) training_stats.report('Loss/signs/fake', gen_logits.sign()) @@ -260,7 +412,9 @@ def accumulate_gradients(self, phase, real_img, real_c, gen_z, gen_c, gain, cur_ real_img_tmp_image_raw = real_img['image_raw'].detach().requires_grad_(phase in ['Dreg', 'Dboth']) real_img_tmp = {'image': real_img_tmp_image, 'image_raw': real_img_tmp_image_raw} - real_logits = self.run_D(real_img_tmp, real_c, blur_sigma=blur_sigma) + # TODO: ADD gen_pc to discriminator + real_logits = self.run_D(real_img_tmp, real_c, blur_sigma=blur_sigma, update_emas=True) + training_stats.report('Loss/scores/real', real_logits) training_stats.report('Loss/signs/real', real_logits.sign()) diff --git a/eg3d/training/networks_stylegan2_pcws.py b/eg3d/training/networks_stylegan2_pcws.py new file mode 100644 index 00000000..e9c74aff --- /dev/null +++ b/eg3d/training/networks_stylegan2_pcws.py @@ -0,0 +1,1041 @@ +# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +"""Network architectures from the paper +"Analyzing and Improving the Image Quality of StyleGAN". +Matches the original implementation of configs E-F by Karras et al. at +https://github.com/NVlabs/stylegan2/blob/master/training/networks_stylegan2.py""" + +from ctypes.wintypes import PCHAR +from re import A +import numpy as np +import torch +from torch_utils import misc +from torch_utils import persistence +from torch_utils.ops import conv2d_resample +from torch_utils.ops import upfirdn2d +from torch_utils.ops import bias_act +from torch_utils.ops import fma + +from ipdb import set_trace as st +import multiprocessing +from torch_utils.utils_ds import grp_range_torch, parallel_FPS +import torch.nn as nn +import torch.nn.functional as F +import torch_scatter +import spconv.pytorch.conv as spconv +from training.costregnet import CostRegNet_Deeper, PcWsUnet +#---------------------------------------------------------------------------- + +@misc.profiled_function +def normalize_2nd_moment(x, dim=1, eps=1e-8): + return x * (x.square().mean(dim=dim, keepdim=True) + eps).rsqrt() + +#---------------------------------------------------------------------------- + +@misc.profiled_function +def modulated_conv2d( + x, # Input tensor of shape [batch_size, in_channels, in_height, in_width]. + weight, # Weight tensor of shape [out_channels, in_channels, kernel_height, kernel_width]. + styles, # Modulation coefficients of shape [batch_size, in_channels]. + noise = None, # Optional noise tensor to add to the output activations. + up = 1, # Integer upsampling factor. + down = 1, # Integer downsampling factor. + padding = 0, # Padding with respect to the upsampled image. + resample_filter = None, # Low-pass filter to apply when resampling activations. Must be prepared beforehand by calling upfirdn2d.setup_filter(). + demodulate = True, # Apply weight demodulation? + flip_weight = True, # False = convolution, True = correlation (matches torch.nn.functional.conv2d). + fused_modconv = True, # Perform modulation, convolution, and demodulation as a single fused operation? +): + batch_size = x.shape[0] + out_channels, in_channels, kh, kw = weight.shape + misc.assert_shape(weight, [out_channels, in_channels, kh, kw]) # [OIkk] + misc.assert_shape(x, [batch_size, in_channels, None, None]) # [NIHW] + misc.assert_shape(styles, [batch_size, in_channels]) # [NI] + + # Pre-normalize inputs to avoid FP16 overflow. + if x.dtype == torch.float16 and demodulate: + weight = weight * (1 / np.sqrt(in_channels * kh * kw) / weight.norm(float('inf'), dim=[1,2,3], keepdim=True)) # max_Ikk + styles = styles / styles.norm(float('inf'), dim=1, keepdim=True) # max_I + + # Calculate per-sample weights and demodulation coefficients. + w = None + dcoefs = None + if demodulate or fused_modconv: + w = weight.unsqueeze(0) # [NOIkk] + w = w * styles.reshape(batch_size, 1, -1, 1, 1) # [NOIkk] + if demodulate: + dcoefs = (w.square().sum(dim=[2,3,4]) + 1e-8).rsqrt() # [NO] + if demodulate and fused_modconv: + w = w * dcoefs.reshape(batch_size, -1, 1, 1, 1) # [NOIkk] + + # Execute by scaling the activations before and after the convolution. + if not fused_modconv: + x = x * styles.to(x.dtype).reshape(batch_size, -1, 1, 1) + x = conv2d_resample.conv2d_resample(x=x, w=weight.to(x.dtype), f=resample_filter, up=up, down=down, padding=padding, flip_weight=flip_weight) + if demodulate and noise is not None: + x = fma.fma(x, dcoefs.to(x.dtype).reshape(batch_size, -1, 1, 1), noise.to(x.dtype)) + elif demodulate: + x = x * dcoefs.to(x.dtype).reshape(batch_size, -1, 1, 1) + elif noise is not None: + x = x.add_(noise.to(x.dtype)) + return x + + # Execute as one fused op using grouped convolution. + with misc.suppress_tracer_warnings(): # this value will be treated as a constant + batch_size = int(batch_size) + misc.assert_shape(x, [batch_size, in_channels, None, None]) + x = x.reshape(1, -1, *x.shape[2:]) + w = w.reshape(-1, in_channels, kh, kw) + x = conv2d_resample.conv2d_resample(x=x, w=w.to(x.dtype), f=resample_filter, up=up, down=down, padding=padding, groups=batch_size, flip_weight=flip_weight) + x = x.reshape(batch_size, -1, *x.shape[2:]) + if noise is not None: + x = x.add_(noise) + return x + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class FullyConnectedLayer(torch.nn.Module): + def __init__(self, + in_features, # Number of input features. + out_features, # Number of output features. + bias = True, # Apply additive bias before the activation function? + activation = 'linear', # Activation function: 'relu', 'lrelu', etc. + lr_multiplier = 1, # Learning rate multiplier. + bias_init = 0, # Initial value for the additive bias. + ): + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.activation = activation + self.weight = torch.nn.Parameter(torch.randn([out_features, in_features]) / lr_multiplier) + self.bias = torch.nn.Parameter(torch.full([out_features], np.float32(bias_init))) if bias else None + self.weight_gain = lr_multiplier / np.sqrt(in_features) + self.bias_gain = lr_multiplier + + def forward(self, x): + + w = self.weight.to(x.dtype) * self.weight_gain + b = self.bias + if b is not None: + b = b.to(x.dtype) + if self.bias_gain != 1: + b = b * self.bias_gain + if self.activation == 'linear' and b is not None: + + x = torch.addmm(b.unsqueeze(0), x, w.t()) + else: + x = x.matmul(w.t()) + x = bias_act.bias_act(x, b, act=self.activation) + return x + + def extra_repr(self): + return f'in_features={self.in_features:d}, out_features={self.out_features:d}, activation={self.activation:s}' + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class Conv2dLayer(torch.nn.Module): + def __init__(self, + in_channels, # Number of input channels. + out_channels, # Number of output channels. + kernel_size, # Width and height of the convolution kernel. + bias = True, # Apply additive bias before the activation function? + activation = 'linear', # Activation function: 'relu', 'lrelu', etc. + up = 1, # Integer upsampling factor. + down = 1, # Integer downsampling factor. + resample_filter = [1,3,3,1], # Low-pass filter to apply when resampling activations. + conv_clamp = None, # Clamp the output to +-X, None = disable clamping. + channels_last = False, # Expect the input to have memory_format=channels_last? + trainable = True, # Update the weights of this layer during training? + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.activation = activation + self.up = up + self.down = down + self.conv_clamp = conv_clamp + self.register_buffer('resample_filter', upfirdn2d.setup_filter(resample_filter)) + self.padding = kernel_size // 2 + self.weight_gain = 1 / np.sqrt(in_channels * (kernel_size ** 2)) + self.act_gain = bias_act.activation_funcs[activation].def_gain + + memory_format = torch.channels_last if channels_last else torch.contiguous_format + weight = torch.randn([out_channels, in_channels, kernel_size, kernel_size]).to(memory_format=memory_format) + bias = torch.zeros([out_channels]) if bias else None + if trainable: + self.weight = torch.nn.Parameter(weight) + self.bias = torch.nn.Parameter(bias) if bias is not None else None + else: + self.register_buffer('weight', weight) + if bias is not None: + self.register_buffer('bias', bias) + else: + self.bias = None + + def forward(self, x, gain=1): + w = self.weight * self.weight_gain + b = self.bias.to(x.dtype) if self.bias is not None else None + flip_weight = (self.up == 1) # slightly faster + x = conv2d_resample.conv2d_resample(x=x, w=w.to(x.dtype), f=self.resample_filter, up=self.up, down=self.down, padding=self.padding, flip_weight=flip_weight) + + act_gain = self.act_gain * gain + act_clamp = self.conv_clamp * gain if self.conv_clamp is not None else None + x = bias_act.bias_act(x, b, act=self.activation, gain=act_gain, clamp=act_clamp) + return x + + def extra_repr(self): + return ' '.join([ + f'in_channels={self.in_channels:d}, out_channels={self.out_channels:d}, activation={self.activation:s},', + f'up={self.up}, down={self.down}']) + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class MappingNetwork(torch.nn.Module): + def __init__(self, + z_dim, # Input latent (Z) dimensionality, 0 = no latent. + c_dim, # Conditioning label (C) dimensionality, 0 = no label. + w_dim, # Intermediate latent (W) dimensionality. + num_ws, # Number of intermediate latents to output, None = do not broadcast. + num_layers = 8, # Number of mapping layers. + embed_features = None, # Label embedding dimensionality, None = same as w_dim. + layer_features = None, # Number of intermediate features in the mapping layers, None = same as w_dim. + activation = 'lrelu', # Activation function: 'relu', 'lrelu', etc. + lr_multiplier = 0.01, # Learning rate multiplier for the mapping layers. + w_avg_beta = 0.998, # Decay for tracking the moving average of W during training, None = do not track. + ): + super().__init__() + self.z_dim = z_dim + self.c_dim = c_dim + self.w_dim = w_dim + self.num_ws = num_ws + self.num_layers = num_layers + self.w_avg_beta = w_avg_beta + + if embed_features is None: + embed_features = w_dim + if c_dim == 0: + embed_features = 0 + if layer_features is None: + layer_features = w_dim + features_list = [z_dim + embed_features] + [layer_features] * (num_layers - 1) + [w_dim] + + if c_dim > 0: + self.embed = FullyConnectedLayer(c_dim, embed_features) + for idx in range(num_layers): + in_features = features_list[idx] + out_features = features_list[idx + 1] + layer = FullyConnectedLayer(in_features, out_features, activation=activation, lr_multiplier=lr_multiplier) + setattr(self, f'fc{idx}', layer) + + if num_ws is not None and w_avg_beta is not None: + self.register_buffer('w_avg', torch.zeros([w_dim])) + + def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, update_emas=False): + # Embed, normalize, and concat inputs. + x = None + with torch.autograd.profiler.record_function('input'): + if self.z_dim > 0: + misc.assert_shape(z, [None, self.z_dim]) + x = normalize_2nd_moment(z.to(torch.float32)) + if self.c_dim > 0: + misc.assert_shape(c, [None, self.c_dim]) + y = normalize_2nd_moment(self.embed(c.to(torch.float32))) + x = torch.cat([x, y], dim=1) if x is not None else y + + # Main layers. + for idx in range(self.num_layers): + layer = getattr(self, f'fc{idx}') + x = layer(x) + # st() # x: (1,512) + + # Update moving average of W. + if update_emas and self.w_avg_beta is not None: + with torch.autograd.profiler.record_function('update_w_avg'): + self.w_avg.copy_(x.detach().mean(dim=0).lerp(self.w_avg, self.w_avg_beta)) + + # Broadcast. + if self.num_ws is not None: + with torch.autograd.profiler.record_function('broadcast'): + x = x.unsqueeze(1).repeat([1, self.num_ws, 1]) + + # Apply truncation. + # st() + if truncation_psi != 1: + with torch.autograd.profiler.record_function('truncate'): + assert self.w_avg_beta is not None + if self.num_ws is None or truncation_cutoff is None: + x = self.w_avg.lerp(x, truncation_psi) + else: + x[:, :truncation_cutoff] = self.w_avg.lerp(x[:, :truncation_cutoff], truncation_psi) + return x + + def extra_repr(self): + return f'z_dim={self.z_dim:d}, c_dim={self.c_dim:d}, w_dim={self.w_dim:d}, num_ws={self.num_ws:d}' + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class SynthesisLayer(torch.nn.Module): + def __init__(self, + in_channels, # Number of input channels. + out_channels, # Number of output channels. + w_dim, # Intermediate latent (W) dimensionality. + resolution, # Resolution of this layer. + kernel_size = 3, # Convolution kernel size. + up = 1, # Integer upsampling factor. + use_noise = True, # Enable noise input? + activation = 'lrelu', # Activation function: 'relu', 'lrelu', etc. + resample_filter = [1,3,3,1], # Low-pass filter to apply when resampling activations. + conv_clamp = None, # Clamp the output of convolution layers to +-X, None = disable clamping. + channels_last = False, # Use channels_last format for the weights? + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.w_dim = w_dim + self.resolution = resolution + self.up = up + self.use_noise = use_noise + self.activation = activation + self.conv_clamp = conv_clamp + self.register_buffer('resample_filter', upfirdn2d.setup_filter(resample_filter)) + self.padding = kernel_size // 2 + self.act_gain = bias_act.activation_funcs[activation].def_gain + + self.affine = FullyConnectedLayer(w_dim, in_channels, bias_init=1) + memory_format = torch.channels_last if channels_last else torch.contiguous_format + self.weight = torch.nn.Parameter(torch.randn([out_channels, in_channels, kernel_size, kernel_size]).to(memory_format=memory_format)) + if use_noise: + self.register_buffer('noise_const', torch.randn([resolution, resolution])) + self.noise_strength = torch.nn.Parameter(torch.zeros([])) + self.bias = torch.nn.Parameter(torch.zeros([out_channels])) + + def forward(self, x, w, noise_mode='random', fused_modconv=True, gain=1): + assert noise_mode in ['random', 'const', 'none'] + in_resolution = self.resolution // self.up + misc.assert_shape(x, [None, self.in_channels, in_resolution, in_resolution]) + styles = self.affine(w) + + noise = None + if self.use_noise and noise_mode == 'random': + noise = torch.randn([x.shape[0], 1, self.resolution, self.resolution], device=x.device) * self.noise_strength + if self.use_noise and noise_mode == 'const': + noise = self.noise_const * self.noise_strength + + flip_weight = (self.up == 1) # slightly faster + x = modulated_conv2d(x=x, weight=self.weight, styles=styles, noise=noise, up=self.up, + padding=self.padding, resample_filter=self.resample_filter, flip_weight=flip_weight, fused_modconv=fused_modconv) + + act_gain = self.act_gain * gain + act_clamp = self.conv_clamp * gain if self.conv_clamp is not None else None + x = bias_act.bias_act(x, self.bias.to(x.dtype), act=self.activation, gain=act_gain, clamp=act_clamp) + return x + + def extra_repr(self): + return ' '.join([ + f'in_channels={self.in_channels:d}, out_channels={self.out_channels:d}, w_dim={self.w_dim:d},', + f'resolution={self.resolution:d}, up={self.up}, activation={self.activation:s}']) + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class ToRGBLayer(torch.nn.Module): + def __init__(self, in_channels, out_channels, w_dim, kernel_size=1, conv_clamp=None, channels_last=False): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.w_dim = w_dim + self.conv_clamp = conv_clamp + self.affine = FullyConnectedLayer(w_dim, in_channels, bias_init=1) + memory_format = torch.channels_last if channels_last else torch.contiguous_format + self.weight = torch.nn.Parameter(torch.randn([out_channels, in_channels, kernel_size, kernel_size]).to(memory_format=memory_format)) + self.bias = torch.nn.Parameter(torch.zeros([out_channels])) + self.weight_gain = 1 / np.sqrt(in_channels * (kernel_size ** 2)) + + def forward(self, x, w, fused_modconv=True): + styles = self.affine(w) * self.weight_gain + x = modulated_conv2d(x=x, weight=self.weight, styles=styles, demodulate=False, fused_modconv=fused_modconv) + x = bias_act.bias_act(x, self.bias.to(x.dtype), clamp=self.conv_clamp) + return x + + def extra_repr(self): + return f'in_channels={self.in_channels:d}, out_channels={self.out_channels:d}, w_dim={self.w_dim:d}' + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class SynthesisBlock(torch.nn.Module): + def __init__(self, + in_channels, # Number of input channels, 0 = first block. + out_channels, # Number of output channels. + w_dim, # Intermediate latent (W) dimensionality. + resolution, # Resolution of this block. + img_channels, # Number of output color channels. + is_last, # Is this the last block? + architecture = 'skip', # Architecture: 'orig', 'skip', 'resnet'. + resample_filter = [1,3,3,1], # Low-pass filter to apply when resampling activations. + conv_clamp = 256, # Clamp the output of convolution layers to +-X, None = disable clamping. + use_fp16 = False, # Use FP16 for this block? + fp16_channels_last = False, # Use channels-last memory format with FP16? + fused_modconv_default = True, # Default value of fused_modconv. 'inference_only' = True for inference, False for training. + **layer_kwargs, # Arguments for SynthesisLayer. + ): + assert architecture in ['orig', 'skip', 'resnet'] + super().__init__() + self.in_channels = in_channels + self.w_dim = w_dim + self.resolution = resolution + self.img_channels = img_channels + self.is_last = is_last + self.architecture = architecture + self.use_fp16 = use_fp16 + self.channels_last = (use_fp16 and fp16_channels_last) + self.fused_modconv_default = fused_modconv_default + self.register_buffer('resample_filter', upfirdn2d.setup_filter(resample_filter)) + self.num_conv = 0 + self.num_torgb = 0 + + if in_channels == 0: + self.const = torch.nn.Parameter(torch.randn([out_channels, resolution, resolution])) + + if in_channels != 0: + self.conv0 = SynthesisLayer(in_channels, out_channels, w_dim=w_dim, resolution=resolution, up=2, + resample_filter=resample_filter, conv_clamp=conv_clamp, channels_last=self.channels_last, **layer_kwargs) + self.num_conv += 1 + + self.conv1 = SynthesisLayer(out_channels, out_channels, w_dim=w_dim, resolution=resolution, + conv_clamp=conv_clamp, channels_last=self.channels_last, **layer_kwargs) + self.num_conv += 1 + + if is_last or architecture == 'skip': + self.torgb = ToRGBLayer(out_channels, img_channels, w_dim=w_dim, + conv_clamp=conv_clamp, channels_last=self.channels_last) + self.num_torgb += 1 + + if in_channels != 0 and architecture == 'resnet': + self.skip = Conv2dLayer(in_channels, out_channels, kernel_size=1, bias=False, up=2, + resample_filter=resample_filter, channels_last=self.channels_last) + + def forward(self, x, img, ws, force_fp32=False, fused_modconv=None, update_emas=False, **layer_kwargs): + _ = update_emas # unused + misc.assert_shape(ws, [None, self.num_conv + self.num_torgb, self.w_dim]) + w_iter = iter(ws.unbind(dim=1)) + if ws.device.type != 'cuda': + force_fp32 = True + dtype = torch.float16 if self.use_fp16 and not force_fp32 else torch.float32 + memory_format = torch.channels_last if self.channels_last and not force_fp32 else torch.contiguous_format + if fused_modconv is None: + fused_modconv = self.fused_modconv_default + if fused_modconv == 'inference_only': + fused_modconv = (not self.training) + + # Input. + if self.in_channels == 0: + x = self.const.to(dtype=dtype, memory_format=memory_format) + x = x.unsqueeze(0).repeat([ws.shape[0], 1, 1, 1]) + else: + misc.assert_shape(x, [None, self.in_channels, self.resolution // 2, self.resolution // 2]) + x = x.to(dtype=dtype, memory_format=memory_format) + + # Main layers. + if self.in_channels == 0: + x = self.conv1(x, next(w_iter), fused_modconv=fused_modconv, **layer_kwargs) + elif self.architecture == 'resnet': + y = self.skip(x, gain=np.sqrt(0.5)) + x = self.conv0(x, next(w_iter), fused_modconv=fused_modconv, **layer_kwargs) + x = self.conv1(x, next(w_iter), fused_modconv=fused_modconv, gain=np.sqrt(0.5), **layer_kwargs) + x = y.add_(x) + else: + x = self.conv0(x, next(w_iter), fused_modconv=fused_modconv, **layer_kwargs) + x = self.conv1(x, next(w_iter), fused_modconv=fused_modconv, **layer_kwargs) + + # ToRGB. + if img is not None: + misc.assert_shape(img, [None, self.img_channels, self.resolution // 2, self.resolution // 2]) + img = upfirdn2d.upsample2d(img, self.resample_filter) + if self.is_last or self.architecture == 'skip': + y = self.torgb(x, next(w_iter), fused_modconv=fused_modconv) + y = y.to(dtype=torch.float32, memory_format=torch.contiguous_format) + img = img.add_(y) if img is not None else y + + assert x.dtype == dtype + assert img is None or img.dtype == torch.float32 + return x, img + + def extra_repr(self): + return f'resolution={self.resolution:d}, architecture={self.architecture:s}' + +#---------------------------------------------------------------------------- + +# directly modify this to generate volume +@persistence.persistent_class +class SynthesisNetwork(torch.nn.Module): + def __init__(self, + w_dim, # Intermediate latent (W) dimensionality. + grid_size, + img_resolution, # Output image resolution. + img_channels, # Number of color channels. + channel_base = 32768, # Overall multiplier for the number of channels. + channel_max = 512, # Maximum number of channels in any layer. + num_fp16_res = 4, # Use FP16 for the N highest resolutions. + **block_kwargs, # Arguments for SynthesisBlock. + ): + assert img_resolution >= 4 and img_resolution & (img_resolution - 1) == 0 + super().__init__() + self.w_dim = w_dim + self.img_resolution = img_resolution + self.img_resolution_log2 = int(np.log2(img_resolution)) + self.img_channels = img_channels + self.num_fp16_res = num_fp16_res + self.block_resolutions = [2 ** i for i in range(2, self.img_resolution_log2 + 1)] + channels_dict = {res: min(channel_base // res, channel_max) for res in self.block_resolutions} + fp16_resolution = max(2 ** (self.img_resolution_log2 + 1 - num_fp16_res), 8) + + self.num_ws = 0 + for res in self.block_resolutions: + in_channels = channels_dict[res // 2] if res > 4 else 0 + out_channels = channels_dict[res] + use_fp16 = (res >= fp16_resolution) + is_last = (res == self.img_resolution) + block = SynthesisBlock(in_channels, out_channels, w_dim=w_dim, resolution=res, + img_channels=img_channels, is_last=is_last, use_fp16=use_fp16, **block_kwargs) + self.num_ws += block.num_conv + if is_last: + self.num_ws += block.num_torgb + setattr(self, f'b{res}', block) + + ###### hard-code attr for voxelize ####### + self.vfe_model = PointNet(fea_dim=9, out_pt_fea_dim=32) # TODO: modify this hard-coded thing + self.pt_selection = 'random' + self.max_pt = 256 + self.pt_pooling = 'max' + if self.pt_pooling == 'max': + self.pool_dim = 64 + self.pos_enc_dim = 0 + # self.fea_compre = 32 + # self.fea_compression = nn.Sequential( + # nn.Linear(self.pool_dim+self.pos_enc_dim, self.fea_compre), + # nn.ReLU() + # ).cuda() + ######### for unet3d ############ + unet_in_channels = 32 + unet_out_channels = 8 + self.grid_res = grid_size + self.grid_size=[self.grid_res]*3 + st() + # self.unet3d=CostRegNet_Deeper(unet_in_channels, out_dim=unet_out_channels, norm_act= nn.BatchNorm3d).to(torch.device("cuda")) + self.pc_ws_unet=PcWsUnet(in_channels=unet_in_channels, in_resolution=self.grid_res, \ + block_resolutions=self.block_resolutions, out_dim=unet_out_channels) + + def forward(self, ws, pc, box_warp, **block_kwargs): + # def forward(self, ws, **block_kwargs): + + RETURN_IMG=True + if RETURN_IMG: + + block_ws = [] + ######## latents ---------------- + + with torch.autograd.profiler.record_function('split_ws'): + misc.assert_shape(ws, [None, self.num_ws, self.w_dim]) + ws = ws.to(torch.float32) + w_idx = 0 + for res in self.block_resolutions: + + block = getattr(self, f'b{res}') + block_ws.append(ws.narrow(1, w_idx, block.num_conv + block.num_torgb)) + w_idx += block.num_conv + + + # 1. extract corresponding res of pointcloud feature: which is smiliar to condtition the ws on pc + # # 1.1 voxelize input pc + + B,_,_=pc.shape + _coor, _feature_3d, density_volume, voxel_size = self.voxelize_spconv_sparse_pointnet( + pc=pc, box_warp=box_warp,grid_size=self.grid_size) + _ret = spconv.SparseConvTensor(_feature_3d, _coor.int(), np.array(self.grid_size), + B) # sp_tensor batch = B*V + _feature_3d = _ret.dense(channels_first = True).contiguous() # [B, C, X, Y, Z], C=32 + # st() + volume_res_features = self.pc_ws_unet(_feature_3d) # a dict of pc feature at different resolution + # st() + + # ----change to all with the same global latent------------ + + # v1: no need to process ws: concat all at the bottleneck + # v2: progressively add latents during upconv + + + ########## generate tri-plane ############## + + x = img = None + # st() # pc.shape + for res, cur_ws in zip(self.block_resolutions, block_ws): + block = getattr(self, f'b{res}') + pc_ws = volume_res_features.get(res) + if pc_ws is not None: + pc_ws = pc_ws.unsqueeze(1).repeat(1, cur_ws.shape[1],1) + # print(pc_ws.shape) + # cur_ws[..., :pc_ws.shape[-1]]=pc_ws: inplace operation not allowed + comb_mask = torch.ones_like(cur_ws) + comb_mask[..., :pc_ws.shape[-1]]*=0 + p1d=(0,cur_ws.shape[-1]-pc_ws.shape[-1]) + pc_ws_pad = F.pad(pc_ws, p1d, 'constant', 0) + cur_ws = comb_mask*cur_ws + (1-comb_mask)*pc_ws_pad + # print(cur_ws.shape) + x, img = block(x, img, cur_ws, **block_kwargs) + # st() # align with img.shape: torch.Size([4, 96, 256, 256]): B,C,H,W + return img + + + # ----change to 3D Unet ------------ + + # # 1. voxelize input pc + _grid_size=[64,64,64] + B,_,_=pc.shape + _coor, _feature_3d, density_volume, voxel_size = self.voxelize_spconv_sparse_pointnet( + pc=pc, box_warp=box_warp,grid_size=self.grid_size) + _ret = spconv.SparseConvTensor(_feature_3d, _coor.int(), np.array(self.grid_size), + B) # sp_tensor batch = B*V + _feature_3d = _ret.dense(channels_first = True).contiguous() # [B, C, X, Y, Z], C=32 + + # voxelize_spconv_sparse_pointnet(self, batch_pcl, grid_size=[], + # batch_bbox=None, pointnet_input=None, pyramid_layer=None): + # # 2. 3D Unet: special: with addtional input ws to concate at the bottle neck so that the upconv part can serve as generator + # self.backbone as in mvsnerf.models + _feature_3d = self.unet3d(_feature_3d.contiguous()) # 3d CONV takes [B, C, X, Y, Z] as input + _feature_3d = _feature_3d.permute(0,1,4,3,2) + # st() + volume = _feature_3d + + return volume + + def extra_repr(self): + return ' '.join([ + f'w_dim={self.w_dim:d}, num_ws={self.num_ws:d},', + f'img_resolution={self.img_resolution:d}, img_channels={self.img_channels:d},', + f'num_fp16_res={self.num_fp16_res:d}']) + + + def voxelize_spconv_sparse_pointnet(self, pc, grid_size=[], + box_warp=None, pointnet_input=None): + ######## parameter alignment ######### + batch_pcl = pc[...,:3].unsqueeze(1) + batch_mtl = pc[...,3:].unsqueeze(1) + device=batch_pcl.device + B,V,_,_ = batch_pcl.shape # torch.Size([4, 1, 1500, 9]) + #check box_warp.shape + batch_bbox = torch.tensor([ + [-box_warp/2, -box_warp/2, -box_warp/2], + [box_warp/2, box_warp/2, box_warp/2] + ],device=batch_pcl.device)[None, None,...].repeat(B,V,1,1) + + pointnet_input='local_xyz' # hard-code + feature ='pointnet' + + ######## function logic ######### + grid_size = torch.tensor(grid_size, device=device) + + ## direct batch voxelization + + batch_voxel_size = (batch_bbox[:,:,1:]-batch_bbox[:,:,:1])/grid_size + voxel_size = batch_voxel_size + batch_xyz_cube_pos = torch.div((batch_pcl-batch_bbox[:,:,:1]), batch_voxel_size, rounding_mode='floor') + if pointnet_input== 'local_xyz': + batch_pcl_local = (batch_pcl - batch_bbox[:,:,:1] - batch_xyz_cube_pos*batch_voxel_size) / batch_voxel_size - 0.5 + batch_pcl_local = torch.cat([batch_pcl_local, batch_mtl], dim=-1) #torch.Size([4, 1, 1500, 9]) + + cat_pt_fea, cat_pt_ind = [], [] + for i_batch in range(len(batch_xyz_cube_pos)): + for i_view in range(len(batch_xyz_cube_pos[i_batch])): + cat_pt_fea.append(batch_pcl_local[i_batch, i_view]) + cat_pt_ind.append(F.pad(batch_xyz_cube_pos[i_batch, i_view],(1,0),'constant',value = i_batch*V+i_view)) + cat_pt_fea = torch.cat(cat_pt_fea,dim=0) + cat_pt_ind = torch.cat(cat_pt_ind,dim = 0) + # st() + else: + raise NotImplemented(False) + + pt_num = cat_pt_ind.shape[0] + # shuffle the data + cur_dev = cat_pt_fea.get_device() + shuffled_ind = torch.randperm(pt_num,device = cur_dev) + cat_pt_fea = cat_pt_fea[shuffled_ind,:] + cat_pt_ind = cat_pt_ind[shuffled_ind,:] + + # unique xy grid index + # st() + unq, unq_inv, unq_cnt = torch.unique(cat_pt_ind,return_inverse=True, return_counts=True, dim=0) + unq = unq.type(torch.int64) + + + # subsample pts + if self.pt_selection == 'random': + grp_ind = grp_range_torch(unq_cnt,cur_dev)[torch.argsort(torch.argsort(unq_inv))] # convert the array that is in the order of grid to the order of cat_pt_feature + remain_ind = grp_ind < self.max_pt # randomly sample max_pt points inside a grid + elif self.pt_selection == 'farthest': + unq_ind = np.split(np.argsort(unq_inv.detach().cpu().numpy()), np.cumsum(unq_cnt.detach().cpu().numpy()[:-1])) + remain_ind = np.zeros((pt_num,),dtype = np.bool) + np_cat_fea = cat_pt_fea.detach().cpu().numpy()[:,:3] + pool_in = [] + for i_inds in unq_ind: + if len(i_inds) > self.max_pt: + pool_in.append((np_cat_fea[i_inds,:],self.max_pt)) + if len(pool_in) > 0: + pool = multiprocessing.Pool(multiprocessing.cpu_count()) + FPS_results = pool.starmap(parallel_FPS, pool_in) + pool.close() + pool.join() + count = 0 + for i_inds in unq_ind: + if len(i_inds) <= self.max_pt: + remain_ind[i_inds] = True + else: + remain_ind[i_inds[FPS_results[count]]] = True + count += 1 + + cat_pt_fea = cat_pt_fea[remain_ind,:] + cat_pt_ind = cat_pt_ind[remain_ind,:] + unq_inv = unq_inv[remain_ind] + unq_cnt = torch.clamp(unq_cnt,max=self.max_pt) + # construct density volume from unqcnt + batch_densities_volumes = unq_cnt[...,None] # add one more dim + + if feature=='embedding': + processed_cat_pt_fea = self.voxel_embed(cat_pt_fea) + elif feature=='pointnet': + processed_cat_pt_fea = self.vfe_model(cat_pt_fea) # global pointnet + # st() + + if self.pt_pooling == 'max': + pooled_data = torch_scatter.scatter_max(processed_cat_pt_fea, unq_inv, dim=0)[0] # choose the max feature for each grid + else: raise NotImplementedError + + # if self.fea_compre: + # processed_pooled_data = self.fea_compression(pooled_data) + # else: + processed_pooled_data = pooled_data + + return unq, processed_pooled_data, batch_densities_volumes, voxel_size + + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class PointNet(torch.nn.Module): + # def __init__(self, cfg): + def __init__(self, fea_dim, out_pt_fea_dim): + super().__init__() + # fea_dim = cfg.DATA_CONFIG.DATALOADER.DATA_DIM + # out_pt_fea_dim = cfg.MODEL.VFE.OUT_CHANNEL + + self.PPmodel = nn.Sequential( + # nn.BatchNorm1d(fea_dim), + nn.Linear(fea_dim, 64), + nn.BatchNorm1d(64), + nn.ReLU(), + nn.Linear(64, 128), + nn.BatchNorm1d(128), + nn.ReLU(), + nn.Linear(128, 256), + nn.BatchNorm1d(256), + nn.ReLU(), + nn.Linear(256, out_pt_fea_dim) + ) + + def forward(self, x): + return self.PPmodel(x) + +#---------------------------------------------------------------------------- + + + +@persistence.persistent_class +class Generator(torch.nn.Module): + def __init__(self, + z_dim, # Input latent (Z) dimensionality. + c_dim, # Conditioning label (C) dimensionality. + w_dim, # Intermediate latent (W) dimensionality. + ####### newly added parameters ###### + pc_dim, # Conditioning poincloud (PC) dimensionality. + volume_res, # Volume resolution. + ########################################## + img_resolution, # Output resolution. + img_channels, # Number of output color channels. + mapping_kwargs = {}, # Arguments for MappingNetwork. + **synthesis_kwargs, # Arguments for SynthesisNetwork. + ): + super().__init__() + self.z_dim = z_dim + self.c_dim = c_dim + self.w_dim = w_dim + ####### newly added parameters ###### + self.pc_dim=pc_dim + self.volume_res=volume_res + ########################################## + self.img_resolution = img_resolution + self.img_channels = img_channels + self.synthesis = SynthesisNetwork(w_dim=w_dim, grid_size=self.volume_res, img_resolution=img_resolution, img_channels=img_channels, **synthesis_kwargs) + self.num_ws = self.synthesis.num_ws + self.mapping = MappingNetwork(z_dim=z_dim, c_dim=c_dim, w_dim=w_dim, num_ws=self.num_ws, **mapping_kwargs) + + def forward(self, z, c, pc, truncation_psi=1, truncation_cutoff=None, update_emas=False, **synthesis_kwargs): + # TODO: whether to include pc info during self.mapping?? + ws = self.mapping(z, c, truncation_psi=truncation_psi, truncation_cutoff=truncation_cutoff, update_emas=update_emas) + img = self.synthesis(ws, pc, update_emas=update_emas, **synthesis_kwargs) + return img + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class DiscriminatorBlock(torch.nn.Module): + def __init__(self, + in_channels, # Number of input channels, 0 = first block. + tmp_channels, # Number of intermediate channels. + out_channels, # Number of output channels. + resolution, # Resolution of this block. + img_channels, # Number of input color channels. + first_layer_idx, # Index of the first layer. + architecture = 'resnet', # Architecture: 'orig', 'skip', 'resnet'. + activation = 'lrelu', # Activation function: 'relu', 'lrelu', etc. + resample_filter = [1,3,3,1], # Low-pass filter to apply when resampling activations. + conv_clamp = None, # Clamp the output of convolution layers to +-X, None = disable clamping. + use_fp16 = False, # Use FP16 for this block? + fp16_channels_last = False, # Use channels-last memory format with FP16? + freeze_layers = 0, # Freeze-D: Number of layers to freeze. + ): + assert in_channels in [0, tmp_channels] + assert architecture in ['orig', 'skip', 'resnet'] + super().__init__() + self.in_channels = in_channels + self.resolution = resolution + self.img_channels = img_channels + self.first_layer_idx = first_layer_idx + self.architecture = architecture + self.use_fp16 = use_fp16 + self.channels_last = (use_fp16 and fp16_channels_last) + self.register_buffer('resample_filter', upfirdn2d.setup_filter(resample_filter)) + + self.num_layers = 0 + def trainable_gen(): + while True: + layer_idx = self.first_layer_idx + self.num_layers + trainable = (layer_idx >= freeze_layers) + self.num_layers += 1 + yield trainable + trainable_iter = trainable_gen() + + if in_channels == 0 or architecture == 'skip': + self.fromrgb = Conv2dLayer(img_channels, tmp_channels, kernel_size=1, activation=activation, + trainable=next(trainable_iter), conv_clamp=conv_clamp, channels_last=self.channels_last) + + self.conv0 = Conv2dLayer(tmp_channels, tmp_channels, kernel_size=3, activation=activation, + trainable=next(trainable_iter), conv_clamp=conv_clamp, channels_last=self.channels_last) + + self.conv1 = Conv2dLayer(tmp_channels, out_channels, kernel_size=3, activation=activation, down=2, + trainable=next(trainable_iter), resample_filter=resample_filter, conv_clamp=conv_clamp, channels_last=self.channels_last) + + if architecture == 'resnet': + self.skip = Conv2dLayer(tmp_channels, out_channels, kernel_size=1, bias=False, down=2, + trainable=next(trainable_iter), resample_filter=resample_filter, channels_last=self.channels_last) + + def forward(self, x, img, force_fp32=False): + if (x if x is not None else img).device.type != 'cuda': + force_fp32 = True + dtype = torch.float16 if self.use_fp16 and not force_fp32 else torch.float32 + memory_format = torch.channels_last if self.channels_last and not force_fp32 else torch.contiguous_format + + # Input. + if x is not None: + misc.assert_shape(x, [None, self.in_channels, self.resolution, self.resolution]) + x = x.to(dtype=dtype, memory_format=memory_format) + + # FromRGB. + if self.in_channels == 0 or self.architecture == 'skip': + misc.assert_shape(img, [None, self.img_channels, self.resolution, self.resolution]) + img = img.to(dtype=dtype, memory_format=memory_format) + y = self.fromrgb(img) + x = x + y if x is not None else y + img = upfirdn2d.downsample2d(img, self.resample_filter) if self.architecture == 'skip' else None + + # Main layers. + if self.architecture == 'resnet': + y = self.skip(x, gain=np.sqrt(0.5)) + x = self.conv0(x) + x = self.conv1(x, gain=np.sqrt(0.5)) + x = y.add_(x) + else: + x = self.conv0(x) + x = self.conv1(x) + + assert x.dtype == dtype + return x, img + + def extra_repr(self): + return f'resolution={self.resolution:d}, architecture={self.architecture:s}' + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class MinibatchStdLayer(torch.nn.Module): + def __init__(self, group_size, num_channels=1): + super().__init__() + self.group_size = group_size + self.num_channels = num_channels + + def forward(self, x): + N, C, H, W = x.shape + with misc.suppress_tracer_warnings(): # as_tensor results are registered as constants + G = torch.min(torch.as_tensor(self.group_size), torch.as_tensor(N)) if self.group_size is not None else N + F = self.num_channels + c = C // F + + y = x.reshape(G, -1, F, c, H, W) # [GnFcHW] Split minibatch N into n groups of size G, and channels C into F groups of size c. + y = y - y.mean(dim=0) # [GnFcHW] Subtract mean over group. + y = y.square().mean(dim=0) # [nFcHW] Calc variance over group. + y = (y + 1e-8).sqrt() # [nFcHW] Calc stddev over group. + y = y.mean(dim=[2,3,4]) # [nF] Take average over channels and pixels. + y = y.reshape(-1, F, 1, 1) # [nF11] Add missing dimensions. + y = y.repeat(G, 1, H, W) # [NFHW] Replicate over group and pixels. + x = torch.cat([x, y], dim=1) # [NCHW] Append to input as new channels. + return x + + def extra_repr(self): + return f'group_size={self.group_size}, num_channels={self.num_channels:d}' + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class DiscriminatorEpilogue(torch.nn.Module): + def __init__(self, + in_channels, # Number of input channels. + cmap_dim, # Dimensionality of mapped conditioning label, 0 = no label. + resolution, # Resolution of this block. + img_channels, # Number of input color channels. + architecture = 'resnet', # Architecture: 'orig', 'skip', 'resnet'. + mbstd_group_size = 4, # Group size for the minibatch standard deviation layer, None = entire minibatch. + mbstd_num_channels = 1, # Number of features for the minibatch standard deviation layer, 0 = disable. + activation = 'lrelu', # Activation function: 'relu', 'lrelu', etc. + conv_clamp = None, # Clamp the output of convolution layers to +-X, None = disable clamping. + ): + assert architecture in ['orig', 'skip', 'resnet'] + super().__init__() + self.in_channels = in_channels + self.cmap_dim = cmap_dim + self.resolution = resolution + self.img_channels = img_channels + self.architecture = architecture + + if architecture == 'skip': + self.fromrgb = Conv2dLayer(img_channels, in_channels, kernel_size=1, activation=activation) + self.mbstd = MinibatchStdLayer(group_size=mbstd_group_size, num_channels=mbstd_num_channels) if mbstd_num_channels > 0 else None + self.conv = Conv2dLayer(in_channels + mbstd_num_channels, in_channels, kernel_size=3, activation=activation, conv_clamp=conv_clamp) + self.fc = FullyConnectedLayer(in_channels * (resolution ** 2), in_channels, activation=activation) + self.out = FullyConnectedLayer(in_channels, 1 if cmap_dim == 0 else cmap_dim) + + def forward(self, x, img, cmap, force_fp32=False): + misc.assert_shape(x, [None, self.in_channels, self.resolution, self.resolution]) # [NCHW] + _ = force_fp32 # unused + dtype = torch.float32 + memory_format = torch.contiguous_format + + # FromRGB. + x = x.to(dtype=dtype, memory_format=memory_format) + if self.architecture == 'skip': + misc.assert_shape(img, [None, self.img_channels, self.resolution, self.resolution]) + img = img.to(dtype=dtype, memory_format=memory_format) + x = x + self.fromrgb(img) + + # Main layers. + if self.mbstd is not None: + x = self.mbstd(x) + x = self.conv(x) + x = self.fc(x.flatten(1)) + x = self.out(x) + + # Conditioning. + if self.cmap_dim > 0: + misc.assert_shape(cmap, [None, self.cmap_dim]) + x = (x * cmap).sum(dim=1, keepdim=True) * (1 / np.sqrt(self.cmap_dim)) + + assert x.dtype == dtype + return x + + def extra_repr(self): + return f'resolution={self.resolution:d}, architecture={self.architecture:s}' + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class Discriminator(torch.nn.Module): + def __init__(self, + c_dim, # Conditioning label (C) dimensionality. + img_resolution, # Input resolution. + img_channels, # Number of input color channels. + architecture = 'resnet', # Architecture: 'orig', 'skip', 'resnet'. + channel_base = 32768, # Overall multiplier for the number of channels. + channel_max = 512, # Maximum number of channels in any layer. + num_fp16_res = 4, # Use FP16 for the N highest resolutions. + conv_clamp = 256, # Clamp the output of convolution layers to +-X, None = disable clamping. + cmap_dim = None, # Dimensionality of mapped conditioning label, None = default. + block_kwargs = {}, # Arguments for DiscriminatorBlock. + mapping_kwargs = {}, # Arguments for MappingNetwork. + epilogue_kwargs = {}, # Arguments for DiscriminatorEpilogue. + ): + super().__init__() + self.c_dim = c_dim + self.img_resolution = img_resolution + self.img_resolution_log2 = int(np.log2(img_resolution)) + self.img_channels = img_channels + self.block_resolutions = [2 ** i for i in range(self.img_resolution_log2, 2, -1)] + channels_dict = {res: min(channel_base // res, channel_max) for res in self.block_resolutions + [4]} + fp16_resolution = max(2 ** (self.img_resolution_log2 + 1 - num_fp16_res), 8) + + if cmap_dim is None: + cmap_dim = channels_dict[4] + if c_dim == 0: + cmap_dim = 0 + + common_kwargs = dict(img_channels=img_channels, architecture=architecture, conv_clamp=conv_clamp) + cur_layer_idx = 0 + for res in self.block_resolutions: + in_channels = channels_dict[res] if res < img_resolution else 0 + tmp_channels = channels_dict[res] + out_channels = channels_dict[res // 2] + use_fp16 = (res >= fp16_resolution) + block = DiscriminatorBlock(in_channels, tmp_channels, out_channels, resolution=res, + first_layer_idx=cur_layer_idx, use_fp16=use_fp16, **block_kwargs, **common_kwargs) + setattr(self, f'b{res}', block) + cur_layer_idx += block.num_layers + if c_dim > 0: + self.mapping = MappingNetwork(z_dim=0, c_dim=c_dim, w_dim=cmap_dim, num_ws=None, w_avg_beta=None, **mapping_kwargs) + self.b4 = DiscriminatorEpilogue(channels_dict[4], cmap_dim=cmap_dim, resolution=4, **epilogue_kwargs, **common_kwargs) + + def forward(self, img, c, update_emas=False, **block_kwargs): + _ = update_emas # unused + x = None + for res in self.block_resolutions: + block = getattr(self, f'b{res}') + x, img = block(x, img, **block_kwargs) + + cmap = None + if self.c_dim > 0: + cmap = self.mapping(None, c) + x = self.b4(x, img, cmap) + return x + + def extra_repr(self): + return f'c_dim={self.c_dim:d}, img_resolution={self.img_resolution:d}, img_channels={self.img_channels:d}' + +#---------------------------------------------------------------------------- \ No newline at end of file diff --git a/eg3d/training/networks_stylegan2_syn_unet.py b/eg3d/training/networks_stylegan2_syn_unet.py new file mode 100644 index 00000000..97f04e73 --- /dev/null +++ b/eg3d/training/networks_stylegan2_syn_unet.py @@ -0,0 +1,1061 @@ +# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +"""Network architectures from the paper +"Analyzing and Improving the Image Quality of StyleGAN". +Matches the original implementation of configs E-F by Karras et al. at +https://github.com/NVlabs/stylegan2/blob/master/training/networks_stylegan2.py""" + +from ctypes.wintypes import PCHAR +from re import A +import numpy as np +import torch +from torch_utils import misc +from torch_utils import persistence +from torch_utils.ops import conv2d_resample +from torch_utils.ops import upfirdn2d +from torch_utils.ops import bias_act +from torch_utils.ops import fma + +from ipdb import set_trace as st +import multiprocessing +from torch_utils.utils_ds import grp_range_torch, parallel_FPS, get_embedder +import torch.nn as nn +import torch.nn.functional as F +import torch_scatter +import spconv.pytorch.conv as spconv +from training.costregnet import CostRegNet_Deeper, Synthesis3DUnet +#---------------------------------------------------------------------------- + +@misc.profiled_function +def normalize_2nd_moment(x, dim=1, eps=1e-8): + return x * (x.square().mean(dim=dim, keepdim=True) + eps).rsqrt() + +#---------------------------------------------------------------------------- + +@misc.profiled_function +def modulated_conv2d( + x, # Input tensor of shape [batch_size, in_channels, in_height, in_width]. + weight, # Weight tensor of shape [out_channels, in_channels, kernel_height, kernel_width]. + styles, # Modulation coefficients of shape [batch_size, in_channels]. + noise = None, # Optional noise tensor to add to the output activations. + up = 1, # Integer upsampling factor. + down = 1, # Integer downsampling factor. + padding = 0, # Padding with respect to the upsampled image. + resample_filter = None, # Low-pass filter to apply when resampling activations. Must be prepared beforehand by calling upfirdn2d.setup_filter(). + demodulate = True, # Apply weight demodulation? + flip_weight = True, # False = convolution, True = correlation (matches torch.nn.functional.conv2d). + fused_modconv = True, # Perform modulation, convolution, and demodulation as a single fused operation? +): + batch_size = x.shape[0] + out_channels, in_channels, kh, kw = weight.shape + misc.assert_shape(weight, [out_channels, in_channels, kh, kw]) # [OIkk] + misc.assert_shape(x, [batch_size, in_channels, None, None]) # [NIHW] + misc.assert_shape(styles, [batch_size, in_channels]) # [NI] + + # Pre-normalize inputs to avoid FP16 overflow. + if x.dtype == torch.float16 and demodulate: + weight = weight * (1 / np.sqrt(in_channels * kh * kw) / weight.norm(float('inf'), dim=[1,2,3], keepdim=True)) # max_Ikk + styles = styles / styles.norm(float('inf'), dim=1, keepdim=True) # max_I + + # Calculate per-sample weights and demodulation coefficients. + w = None + dcoefs = None + st() + if demodulate or fused_modconv: + w = weight.unsqueeze(0) # [NOIkk] + w = w * styles.reshape(batch_size, 1, -1, 1, 1) # [NOIkk] + if demodulate: + dcoefs = (w.square().sum(dim=[2,3,4]) + 1e-8).rsqrt() # [NO] + if demodulate and fused_modconv: + w = w * dcoefs.reshape(batch_size, -1, 1, 1, 1) # [NOIkk] + + # Execute by scaling the activations before and after the convolution. + if not fused_modconv: + x = x * styles.to(x.dtype).reshape(batch_size, -1, 1, 1) + x = conv2d_resample.conv2d_resample(x=x, w=weight.to(x.dtype), f=resample_filter, up=up, down=down, padding=padding, flip_weight=flip_weight) + if demodulate and noise is not None: + x = fma.fma(x, dcoefs.to(x.dtype).reshape(batch_size, -1, 1, 1), noise.to(x.dtype)) + elif demodulate: + x = x * dcoefs.to(x.dtype).reshape(batch_size, -1, 1, 1) + elif noise is not None: + x = x.add_(noise.to(x.dtype)) + return x + + # Execute as one fused op using grouped convolution. + with misc.suppress_tracer_warnings(): # this value will be treated as a constant + batch_size = int(batch_size) + misc.assert_shape(x, [batch_size, in_channels, None, None]) + x = x.reshape(1, -1, *x.shape[2:]) + w = w.reshape(-1, in_channels, kh, kw) + x = conv2d_resample.conv2d_resample(x=x, w=w.to(x.dtype), f=resample_filter, up=up, down=down, padding=padding, groups=batch_size, flip_weight=flip_weight) + x = x.reshape(batch_size, -1, *x.shape[2:]) + if noise is not None: + # st() + x = x.add_(noise) + return x + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class FullyConnectedLayer(torch.nn.Module): + def __init__(self, + in_features, # Number of input features. + out_features, # Number of output features. + bias = True, # Apply additive bias before the activation function? + activation = 'linear', # Activation function: 'relu', 'lrelu', etc. + lr_multiplier = 1, # Learning rate multiplier. + bias_init = 0, # Initial value for the additive bias. + ): + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.activation = activation + self.weight = torch.nn.Parameter(torch.randn([out_features, in_features]) / lr_multiplier) + self.bias = torch.nn.Parameter(torch.full([out_features], np.float32(bias_init))) if bias else None + self.weight_gain = lr_multiplier / np.sqrt(in_features) + self.bias_gain = lr_multiplier + + def forward(self, x): + + w = self.weight.to(x.dtype) * self.weight_gain + b = self.bias + if b is not None: + b = b.to(x.dtype) + if self.bias_gain != 1: + b = b * self.bias_gain + if self.activation == 'linear' and b is not None: + + x = torch.addmm(b.unsqueeze(0), x, w.t()) + else: + x = x.matmul(w.t()) + x = bias_act.bias_act(x, b, act=self.activation) + return x + + def extra_repr(self): + return f'in_features={self.in_features:d}, out_features={self.out_features:d}, activation={self.activation:s}' + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class Conv2dLayer(torch.nn.Module): + def __init__(self, + in_channels, # Number of input channels. + out_channels, # Number of output channels. + kernel_size, # Width and height of the convolution kernel. + bias = True, # Apply additive bias before the activation function? + activation = 'linear', # Activation function: 'relu', 'lrelu', etc. + up = 1, # Integer upsampling factor. + down = 1, # Integer downsampling factor. + resample_filter = [1,3,3,1], # Low-pass filter to apply when resampling activations. + conv_clamp = None, # Clamp the output to +-X, None = disable clamping. + channels_last = False, # Expect the input to have memory_format=channels_last? + trainable = True, # Update the weights of this layer during training? + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.activation = activation + self.up = up + self.down = down + self.conv_clamp = conv_clamp + self.register_buffer('resample_filter', upfirdn2d.setup_filter(resample_filter)) + self.padding = kernel_size // 2 + self.weight_gain = 1 / np.sqrt(in_channels * (kernel_size ** 2)) + self.act_gain = bias_act.activation_funcs[activation].def_gain + + memory_format = torch.channels_last if channels_last else torch.contiguous_format + weight = torch.randn([out_channels, in_channels, kernel_size, kernel_size]).to(memory_format=memory_format) + bias = torch.zeros([out_channels]) if bias else None + if trainable: + self.weight = torch.nn.Parameter(weight) + self.bias = torch.nn.Parameter(bias) if bias is not None else None + else: + self.register_buffer('weight', weight) + if bias is not None: + self.register_buffer('bias', bias) + else: + self.bias = None + + def forward(self, x, gain=1): + w = self.weight * self.weight_gain + b = self.bias.to(x.dtype) if self.bias is not None else None + flip_weight = (self.up == 1) # slightly faster + x = conv2d_resample.conv2d_resample(x=x, w=w.to(x.dtype), f=self.resample_filter, up=self.up, down=self.down, padding=self.padding, flip_weight=flip_weight) + + act_gain = self.act_gain * gain + act_clamp = self.conv_clamp * gain if self.conv_clamp is not None else None + x = bias_act.bias_act(x, b, act=self.activation, gain=act_gain, clamp=act_clamp) + return x + + def extra_repr(self): + return ' '.join([ + f'in_channels={self.in_channels:d}, out_channels={self.out_channels:d}, activation={self.activation:s},', + f'up={self.up}, down={self.down}']) + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class MappingNetwork(torch.nn.Module): + def __init__(self, + z_dim, # Input latent (Z) dimensionality, 0 = no latent. + c_dim, # Conditioning label (C) dimensionality, 0 = no label. + w_dim, # Intermediate latent (W) dimensionality. + num_ws, # Number of intermediate latents to output, None = do not broadcast. + num_layers = 8, # Number of mapping layers. + embed_features = None, # Label embedding dimensionality, None = same as w_dim. + layer_features = None, # Number of intermediate features in the mapping layers, None = same as w_dim. + activation = 'lrelu', # Activation function: 'relu', 'lrelu', etc. + lr_multiplier = 0.01, # Learning rate multiplier for the mapping layers. + w_avg_beta = 0.998, # Decay for tracking the moving average of W during training, None = do not track. + ): + super().__init__() + self.z_dim = z_dim + self.c_dim = c_dim + self.w_dim = w_dim + self.num_ws = num_ws + self.num_layers = num_layers + self.w_avg_beta = w_avg_beta + + if embed_features is None: + embed_features = w_dim + if c_dim == 0: + embed_features = 0 + if layer_features is None: + layer_features = w_dim + features_list = [z_dim + embed_features] + [layer_features] * (num_layers - 1) + [w_dim] + + if c_dim > 0: + self.embed = FullyConnectedLayer(c_dim, embed_features) + for idx in range(num_layers): + in_features = features_list[idx] + out_features = features_list[idx + 1] + layer = FullyConnectedLayer(in_features, out_features, activation=activation, lr_multiplier=lr_multiplier) + setattr(self, f'fc{idx}', layer) + + if num_ws is not None and w_avg_beta is not None: + self.register_buffer('w_avg', torch.zeros([w_dim])) + + def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, update_emas=False): + # Embed, normalize, and concat inputs. + x = None + with torch.autograd.profiler.record_function('input'): + if self.z_dim > 0: + misc.assert_shape(z, [None, self.z_dim]) + x = normalize_2nd_moment(z.to(torch.float32)) + if self.c_dim > 0: + misc.assert_shape(c, [None, self.c_dim]) + y = normalize_2nd_moment(self.embed(c.to(torch.float32))) + x = torch.cat([x, y], dim=1) if x is not None else y + + # Main layers. + for idx in range(self.num_layers): + layer = getattr(self, f'fc{idx}') + x = layer(x) + + # Update moving average of W. + if update_emas and self.w_avg_beta is not None: + with torch.autograd.profiler.record_function('update_w_avg'): + self.w_avg.copy_(x.detach().mean(dim=0).lerp(self.w_avg, self.w_avg_beta)) + + # Broadcast. + if self.num_ws is not None: + with torch.autograd.profiler.record_function('broadcast'): + x = x.unsqueeze(1).repeat([1, self.num_ws, 1]) + + # Apply truncation. + if truncation_psi != 1: + with torch.autograd.profiler.record_function('truncate'): + assert self.w_avg_beta is not None + if self.num_ws is None or truncation_cutoff is None: + x = self.w_avg.lerp(x, truncation_psi) + else: + x[:, :truncation_cutoff] = self.w_avg.lerp(x[:, :truncation_cutoff], truncation_psi) + return x + + def extra_repr(self): + return f'z_dim={self.z_dim:d}, c_dim={self.c_dim:d}, w_dim={self.w_dim:d}, num_ws={self.num_ws:d}' + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class SynthesisLayer(torch.nn.Module): + def __init__(self, + in_channels, # Number of input channels. + out_channels, # Number of output channels. + w_dim, # Intermediate latent (W) dimensionality. + resolution, # Resolution of this layer. + kernel_size = 3, # Convolution kernel size. + up = 1, # Integer upsampling factor. + use_noise = True, # Enable noise input? + activation = 'lrelu', # Activation function: 'relu', 'lrelu', etc. + resample_filter = [1,3,3,1], # Low-pass filter to apply when resampling activations. + conv_clamp = None, # Clamp the output of convolution layers to +-X, None = disable clamping. + channels_last = False, # Use channels_last format for the weights? + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.w_dim = w_dim + self.resolution = resolution + self.up = up + self.use_noise = use_noise + self.activation = activation + self.conv_clamp = conv_clamp + self.register_buffer('resample_filter', upfirdn2d.setup_filter(resample_filter)) + self.padding = kernel_size // 2 + self.act_gain = bias_act.activation_funcs[activation].def_gain + + self.affine = FullyConnectedLayer(w_dim, in_channels, bias_init=1) + memory_format = torch.channels_last if channels_last else torch.contiguous_format + self.weight = torch.nn.Parameter(torch.randn([out_channels, in_channels, kernel_size, kernel_size]).to(memory_format=memory_format)) + if use_noise: + self.register_buffer('noise_const', torch.randn([resolution, resolution])) + self.noise_strength = torch.nn.Parameter(torch.zeros([])) + self.bias = torch.nn.Parameter(torch.zeros([out_channels])) + + def forward(self, x, w, noise_mode='random', fused_modconv=True, gain=1): + assert noise_mode in ['random', 'const', 'none'] + in_resolution = self.resolution // self.up + misc.assert_shape(x, [None, self.in_channels, in_resolution, in_resolution]) + styles = self.affine(w) + + noise = None + + if self.use_noise and noise_mode == 'random': + noise = torch.randn([x.shape[0], 1, self.resolution, self.resolution], device=x.device) * self.noise_strength + if self.use_noise and noise_mode == 'const': + noise = self.noise_const * self.noise_strength + + flip_weight = (self.up == 1) # slightly faster + x = modulated_conv2d(x=x, weight=self.weight, styles=styles, noise=noise, up=self.up, + padding=self.padding, resample_filter=self.resample_filter, flip_weight=flip_weight, fused_modconv=fused_modconv) + + act_gain = self.act_gain * gain + act_clamp = self.conv_clamp * gain if self.conv_clamp is not None else None + x = bias_act.bias_act(x, self.bias.to(x.dtype), act=self.activation, gain=act_gain, clamp=act_clamp) + return x + + def extra_repr(self): + return ' '.join([ + f'in_channels={self.in_channels:d}, out_channels={self.out_channels:d}, w_dim={self.w_dim:d},', + f'resolution={self.resolution:d}, up={self.up}, activation={self.activation:s}']) + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class ToRGBLayer(torch.nn.Module): + def __init__(self, in_channels, out_channels, w_dim, kernel_size=1, conv_clamp=None, channels_last=False): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.w_dim = w_dim + self.conv_clamp = conv_clamp + self.affine = FullyConnectedLayer(w_dim, in_channels, bias_init=1) + memory_format = torch.channels_last if channels_last else torch.contiguous_format + self.weight = torch.nn.Parameter(torch.randn([out_channels, in_channels, kernel_size, kernel_size]).to(memory_format=memory_format)) + self.bias = torch.nn.Parameter(torch.zeros([out_channels])) + self.weight_gain = 1 / np.sqrt(in_channels * (kernel_size ** 2)) + + def forward(self, x, w, fused_modconv=True): + styles = self.affine(w) * self.weight_gain + x = modulated_conv2d(x=x, weight=self.weight, styles=styles, demodulate=False, fused_modconv=fused_modconv) + x = bias_act.bias_act(x, self.bias.to(x.dtype), clamp=self.conv_clamp) + return x + + def extra_repr(self): + return f'in_channels={self.in_channels:d}, out_channels={self.out_channels:d}, w_dim={self.w_dim:d}' + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class SynthesisBlock(torch.nn.Module): + def __init__(self, + in_channels, # Number of input channels, 0 = first block. + out_channels, # Number of output channels. + w_dim, # Intermediate latent (W) dimensionality. + resolution, # Resolution of this block. + img_channels, # Number of output color channels. + is_last, # Is this the last block? + architecture = 'skip', # Architecture: 'orig', 'skip', 'resnet'. + resample_filter = [1,3,3,1], # Low-pass filter to apply when resampling activations. + conv_clamp = 256, # Clamp the output of convolution layers to +-X, None = disable clamping. + use_fp16 = False, # Use FP16 for this block? + fp16_channels_last = False, # Use channels-last memory format with FP16? + fused_modconv_default = True, # Default value of fused_modconv. 'inference_only' = True for inference, False for training. + **layer_kwargs, # Arguments for SynthesisLayer. + ): + assert architecture in ['orig', 'skip', 'resnet'] + super().__init__() + self.in_channels = in_channels + self.w_dim = w_dim + self.resolution = resolution + self.img_channels = img_channels + self.is_last = is_last + self.architecture = architecture + self.use_fp16 = use_fp16 + self.channels_last = (use_fp16 and fp16_channels_last) + self.fused_modconv_default = fused_modconv_default + self.register_buffer('resample_filter', upfirdn2d.setup_filter(resample_filter)) + self.num_conv = 0 + self.num_torgb = 0 + + if in_channels == 0: + self.const = torch.nn.Parameter(torch.randn([out_channels, resolution, resolution])) + + if in_channels != 0: + self.conv0 = SynthesisLayer(in_channels, out_channels, w_dim=w_dim, resolution=resolution, up=2, + resample_filter=resample_filter, conv_clamp=conv_clamp, channels_last=self.channels_last, **layer_kwargs) + self.num_conv += 1 + + self.conv1 = SynthesisLayer(out_channels, out_channels, w_dim=w_dim, resolution=resolution, + conv_clamp=conv_clamp, channels_last=self.channels_last, **layer_kwargs) + self.num_conv += 1 + + if is_last or architecture == 'skip': + self.torgb = ToRGBLayer(out_channels, img_channels, w_dim=w_dim, + conv_clamp=conv_clamp, channels_last=self.channels_last) + self.num_torgb += 1 + + if in_channels != 0 and architecture == 'resnet': + self.skip = Conv2dLayer(in_channels, out_channels, kernel_size=1, bias=False, up=2, + resample_filter=resample_filter, channels_last=self.channels_last) + + def forward(self, x, img, ws, force_fp32=False, fused_modconv=None, update_emas=False, **layer_kwargs): + _ = update_emas # unused + misc.assert_shape(ws, [None, self.num_conv + self.num_torgb, self.w_dim]) + w_iter = iter(ws.unbind(dim=1)) + if ws.device.type != 'cuda': + force_fp32 = True + dtype = torch.float16 if self.use_fp16 and not force_fp32 else torch.float32 + memory_format = torch.channels_last if self.channels_last and not force_fp32 else torch.contiguous_format + if fused_modconv is None: + fused_modconv = self.fused_modconv_default + if fused_modconv == 'inference_only': + fused_modconv = (not self.training) + + # Input. + if self.in_channels == 0: + x = self.const.to(dtype=dtype, memory_format=memory_format) + x = x.unsqueeze(0).repeat([ws.shape[0], 1, 1, 1]) + else: + misc.assert_shape(x, [None, self.in_channels, self.resolution // 2, self.resolution // 2]) + x = x.to(dtype=dtype, memory_format=memory_format) + + # Main layers. + if self.in_channels == 0: + x = self.conv1(x, next(w_iter), fused_modconv=fused_modconv, **layer_kwargs) + elif self.architecture == 'resnet': + y = self.skip(x, gain=np.sqrt(0.5)) + x = self.conv0(x, next(w_iter), fused_modconv=fused_modconv, **layer_kwargs) + x = self.conv1(x, next(w_iter), fused_modconv=fused_modconv, gain=np.sqrt(0.5), **layer_kwargs) + x = y.add_(x) + else: + x = self.conv0(x, next(w_iter), fused_modconv=fused_modconv, **layer_kwargs) + x = self.conv1(x, next(w_iter), fused_modconv=fused_modconv, **layer_kwargs) + + # ToRGB. + if img is not None: + misc.assert_shape(img, [None, self.img_channels, self.resolution // 2, self.resolution // 2]) + img = upfirdn2d.upsample2d(img, self.resample_filter) + if self.is_last or self.architecture == 'skip': + y = self.torgb(x, next(w_iter), fused_modconv=fused_modconv) + y = y.to(dtype=torch.float32, memory_format=torch.contiguous_format) + img = img.add_(y) if img is not None else y + + assert x.dtype == dtype + assert img is None or img.dtype == torch.float32 + return x, img + + def extra_repr(self): + return f'resolution={self.resolution:d}, architecture={self.architecture:s}' + +#---------------------------------------------------------------------------- + +# directly modify this to generate volume +@persistence.persistent_class +class SynthesisNetwork(torch.nn.Module): + def __init__(self, + w_dim, # Intermediate latent (W) dimensionality. + volume_res, + noise_strength, + # vfe_feature, + img_resolution, # Output image resolution. + img_channels, # Number of color channels. + channel_base = 32768, # Overall multiplier for the number of channels. + channel_max = 512, # Maximum number of channels in any layer. + num_fp16_res = 4, # Use FP16 for the N highest resolutions. + **block_kwargs, # Arguments for SynthesisBlock. + ): + assert img_resolution >= 4 and img_resolution & (img_resolution - 1) == 0 + super().__init__() + self.w_dim = w_dim + self.img_resolution = img_resolution + self.img_resolution_log2 = int(np.log2(img_resolution)) + self.img_channels = img_channels + self.num_fp16_res = num_fp16_res + self.block_resolutions = [2 ** i for i in range(2, self.img_resolution_log2 + 1)] + channels_dict = {res: min(channel_base // res, channel_max) for res in self.block_resolutions} + fp16_resolution = max(2 ** (self.img_resolution_log2 + 1 - num_fp16_res), 8) + + self.num_ws = 0 + for res in self.block_resolutions: + in_channels = channels_dict[res // 2] if res > 4 else 0 + out_channels = channels_dict[res] + use_fp16 = (res >= fp16_resolution) + is_last = (res == self.img_resolution) + block = SynthesisBlock(in_channels, out_channels, w_dim=w_dim, resolution=res, + img_channels=img_channels, is_last=is_last, use_fp16=use_fp16, **block_kwargs) + self.num_ws += block.num_conv + if is_last: + self.num_ws += block.num_torgb + setattr(self, f'b{res}', block) + + ###### hard-code attr for voxelize ####### + # vfe_feature = 'embedding' + vfe_feature = 'pointnet' + self.vfe_feature = vfe_feature + if self.vfe_feature=='embedding': + assert False + # processed_cat_pt_fea = self.voxel_embed(cat_pt_fea) + # EMBEDDER: + embedder_kwargs = { + 'multires': 5, + 'i': 0, # i_embed + 'input_dims': 9, + } + embed_fn, input_ch = get_embedder( + **embedder_kwargs + ) + + # embed_fn, input_ch = get_embedder( + # cfg.MODEL.EMBEDDER.multires, + # cfg.MODEL.EMBEDDER.i_embed, + # input_dims=cfg.MODEL.EMBEDDER.pts_dim) + # input_ch: 33 + self.vfe_model = embed_fn + self.fea_compre = True + self.fea_compression = nn.Sequential( + nn.Linear(input_ch, 32), + nn.ReLU() + ).cuda() + + elif self.vfe_feature=='pointnet': + self.vfe_model = PointNet(fea_dim=9, out_pt_fea_dim=32) # TODO: modify this hard-coded thing + self.fea_compre = False + + self.pt_selection = 'random' + self.max_pt = 256 + self.pt_pooling = 'max' + if self.pt_pooling == 'max': + self.pool_dim = 64 + self.pos_enc_dim = 0 + + ######### for unet3d ############ + self.grid_size=[volume_res]*3 + unet_in_channels = 32 + # self.unet3d=CostRegNet_Deeper(unet_in_channels, norm_act= nn.BatchNorm3d).to(torch.device("cuda")) + self.synthesis_unet3d=Synthesis3DUnet(unet_in_channels, + use_noise=True, noise_strength = noise_strength, norm_act= nn.BatchNorm3d).to(torch.device("cuda")) + + def forward(self, ws, pc, box_warp, **block_kwargs): + # def forward(self, ws, **block_kwargs): + RETURN_IMG=False + if RETURN_IMG: + # st() + block_ws = [] + ######## latents ---------------- + + with torch.autograd.profiler.record_function('split_ws'): + misc.assert_shape(ws, [None, self.num_ws, self.w_dim]) + ws = ws.to(torch.float32) + w_idx = 0 + for res in self.block_resolutions: + block = getattr(self, f'b{res}') + block_ws.append(ws.narrow(1, w_idx, block.num_conv + block.num_torgb)) + w_idx += block.num_conv + + + # ----change to all with the same global latent------------ + + # v1: no need to process ws: concat all at the bottleneck + # v2: progressively add latents during upconv + + + ########## generate tri-plane ############## + + x = img = None + # st() # pc.shape + for res, cur_ws in zip(self.block_resolutions, block_ws): + block = getattr(self, f'b{res}') + x, img = block(x, img, cur_ws, **block_kwargs) + # st() # align with img.shape: torch.Size([4, 96, 256, 256]): B,C,H,W + # target 3d img shape: 1, 32, 64, 64, 64 + return img + + # ----change to 3D Unet ------------ + + # # 1. voxelize input pc + + B,_,_=pc.shape + _coor, _feature_3d, density_volume, voxel_size = self.voxelize_spconv_sparse_pointnet( + pc=pc, box_warp=box_warp,grid_size=self.grid_size) + _ret = spconv.SparseConvTensor(_feature_3d, _coor.int(), np.array(self.grid_size), + B) # sp_tensor batch = B*V + _feature_3d = _ret.dense(channels_first = True).contiguous() # [B, C, X, Y, Z], C=32 + + + # # 2. 3D Unet: special: with addtional input ws to concate at the bottle neck so that the upconv part can serve as generator + + # 2.1 _feature_3d = self.unet3d(_feature_3d.contiguous()) # 3d CONV takes [B, C, X, Y, Z] as input + # 2.2 add latent and noises + _feature_3d = self.synthesis_unet3d(_feature_3d, ws) + + # st() + volume = _feature_3d.permute(0,1,4,3,2) + # st() + + return volume + + def extra_repr(self): + return ' '.join([ + f'w_dim={self.w_dim:d}, num_ws={self.num_ws:d},', + f'img_resolution={self.img_resolution:d}, img_channels={self.img_channels:d},', + f'num_fp16_res={self.num_fp16_res:d}']) + + + def voxelize_spconv_sparse_pointnet(self, pc, grid_size=[], + box_warp=None, pointnet_input=None): + ######## parameter alignment ######### + batch_pcl = pc[...,:3].unsqueeze(1) + batch_mtl = pc[...,3:].unsqueeze(1) + device=batch_pcl.device + B,V,_,_ = batch_pcl.shape # torch.Size([4, 1, 1500, 9]) + #check box_warp.shape + batch_bbox = torch.tensor([ + [-box_warp/2, -box_warp/2, -box_warp/2], + [box_warp/2, box_warp/2, box_warp/2] + ],device=batch_pcl.device)[None, None,...].repeat(B,V,1,1) + + pointnet_input='local_xyz' # hard-code + # feature ='pointnet' + + ######## function logic ######### + + grid_size = torch.tensor(grid_size, device=device) + + ## direct batch voxelization + + batch_voxel_size = (batch_bbox[:,:,1:]-batch_bbox[:,:,:1])/grid_size + voxel_size = batch_voxel_size + batch_xyz_cube_pos = torch.div((batch_pcl-batch_bbox[:,:,:1]), batch_voxel_size, rounding_mode='floor') + if pointnet_input== 'local_xyz': + batch_pcl_local = (batch_pcl - batch_bbox[:,:,:1] - batch_xyz_cube_pos*batch_voxel_size) / batch_voxel_size - 0.5 + batch_pcl_local = torch.cat([batch_pcl_local, batch_mtl], dim=-1) #torch.Size([4, 1, 1500, 9]) + + cat_pt_fea, cat_pt_ind = [], [] + for i_batch in range(len(batch_xyz_cube_pos)): + for i_view in range(len(batch_xyz_cube_pos[i_batch])): + cat_pt_fea.append(batch_pcl_local[i_batch, i_view]) + cat_pt_ind.append(F.pad(batch_xyz_cube_pos[i_batch, i_view],(1,0),'constant',value = i_batch*V+i_view)) + cat_pt_fea = torch.cat(cat_pt_fea,dim=0) + cat_pt_ind = torch.cat(cat_pt_ind,dim = 0) + # st() + else: + raise NotImplemented(False) + + pt_num = cat_pt_ind.shape[0] + # shuffle the data + + cur_dev = cat_pt_fea.get_device() + + shuffled_ind = torch.randperm(pt_num,device = cur_dev) + cat_pt_fea = cat_pt_fea[shuffled_ind,:] + cat_pt_ind = cat_pt_ind[shuffled_ind,:] + + # unique xy grid index + unq, unq_inv, unq_cnt = torch.unique(cat_pt_ind,return_inverse=True, return_counts=True, dim=0) + unq = unq.type(torch.int64) + + + # subsample pts + if self.pt_selection == 'random': + grp_ind = grp_range_torch(unq_cnt,cur_dev)[torch.argsort(torch.argsort(unq_inv))] # convert the array that is in the order of grid to the order of cat_pt_feature + remain_ind = grp_ind < self.max_pt # randomly sample max_pt points inside a grid + elif self.pt_selection == 'farthest': + unq_ind = np.split(np.argsort(unq_inv.detach().cpu().numpy()), np.cumsum(unq_cnt.detach().cpu().numpy()[:-1])) + remain_ind = np.zeros((pt_num,),dtype = np.bool) + np_cat_fea = cat_pt_fea.detach().cpu().numpy()[:,:3] + pool_in = [] + for i_inds in unq_ind: + if len(i_inds) > self.max_pt: + pool_in.append((np_cat_fea[i_inds,:],self.max_pt)) + if len(pool_in) > 0: + pool = multiprocessing.Pool(multiprocessing.cpu_count()) + FPS_results = pool.starmap(parallel_FPS, pool_in) + pool.close() + pool.join() + count = 0 + for i_inds in unq_ind: + if len(i_inds) <= self.max_pt: + remain_ind[i_inds] = True + else: + remain_ind[i_inds[FPS_results[count]]] = True + count += 1 + + cat_pt_fea = cat_pt_fea[remain_ind,:] + cat_pt_ind = cat_pt_ind[remain_ind,:] + unq_inv = unq_inv[remain_ind] + unq_cnt = torch.clamp(unq_cnt,max=self.max_pt) + # construct density volume from unqcnt + batch_densities_volumes = unq_cnt[...,None] # add one more dim + + # if feature=='embedding': + # processed_cat_pt_fea = self.voxel_embed(cat_pt_fea) + # elif feature=='pointnet': + processed_cat_pt_fea = self.vfe_model(cat_pt_fea) # global pointnet + # st() + + if self.pt_pooling == 'max': + pooled_data = torch_scatter.scatter_max(processed_cat_pt_fea, unq_inv, dim=0)[0] # choose the max feature for each grid + else: raise NotImplementedError + + if self.fea_compre: + processed_pooled_data = self.fea_compression(pooled_data) + else: + processed_pooled_data = pooled_data + + return unq, processed_pooled_data, batch_densities_volumes, voxel_size + + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class PointNet(torch.nn.Module): + # def __init__(self, cfg): + def __init__(self, fea_dim, out_pt_fea_dim): + super().__init__() + # fea_dim = cfg.DATA_CONFIG.DATALOADER.DATA_DIM + # out_pt_fea_dim = cfg.MODEL.VFE.OUT_CHANNEL + + # self.PPmodel = nn.Sequential( + # # nn.BatchNorm1d(fea_dim), + # nn.Linear(fea_dim, 64), + # nn.SyncBatchNorm(64), + # nn.ReLU(), + # nn.Linear(64, 128), + # nn.SyncBatchNorm(128), + # nn.ReLU(), + # nn.Linear(128, 256), + # nn.SyncBatchNorm(256), + # nn.ReLU(), + # nn.Linear(256, out_pt_fea_dim) + # ) + + self.PPmodel = nn.Sequential( + # nn.BatchNorm1d(fea_dim), + nn.Linear(fea_dim, 64), + nn.BatchNorm1d(64), + nn.ReLU(), + nn.Linear(64, 128), + nn.BatchNorm1d(128), + nn.ReLU(), + nn.Linear(128, 256), + nn.BatchNorm1d(256), + nn.ReLU(), + nn.Linear(256, out_pt_fea_dim) + ) + + def forward(self, x): + return self.PPmodel(x) + +#---------------------------------------------------------------------------- + + + +@persistence.persistent_class +class Generator(torch.nn.Module): + def __init__(self, + z_dim, # Input latent (Z) dimensionality. + c_dim, # Conditioning label (C) dimensionality. + w_dim, # Intermediate latent (W) dimensionality. + ####### newly added parameters ###### + pc_dim, # Conditioning poincloud (PC) dimensionality. + volume_res, # Volume resolution. + noise_strength, # Factor to multiply with noise in the 3D Unet block. + ########################################## + img_resolution, # Output resolution. + img_channels, # Number of output color channels. + mapping_kwargs = {}, # Arguments for MappingNetwork. + **synthesis_kwargs, # Arguments for SynthesisNetwork. + ): + super().__init__() + self.z_dim = z_dim + self.c_dim = c_dim + self.w_dim = w_dim + ####### newly added parameters ###### + self.pc_dim=pc_dim + self.volume_res=volume_res + ########################################## + self.img_resolution = img_resolution + self.img_channels = img_channels + self.synthesis = SynthesisNetwork(w_dim=w_dim,volume_res=volume_res, img_resolution=img_resolution, noise_strength=noise_strength, img_channels=img_channels, **synthesis_kwargs) + self.num_ws = self.synthesis.num_ws + self.mapping = MappingNetwork(z_dim=z_dim, c_dim=c_dim, w_dim=w_dim, num_ws=self.num_ws, **mapping_kwargs) + + def forward(self, z, c, pc, truncation_psi=1, truncation_cutoff=None, update_emas=False, **synthesis_kwargs): + # TODO: whether to include pc info during self.mapping?? + ws = self.mapping(z, c, truncation_psi=truncation_psi, truncation_cutoff=truncation_cutoff, update_emas=update_emas) + img = self.synthesis(ws, pc, update_emas=update_emas, **synthesis_kwargs) + return img + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class DiscriminatorBlock(torch.nn.Module): + def __init__(self, + in_channels, # Number of input channels, 0 = first block. + tmp_channels, # Number of intermediate channels. + out_channels, # Number of output channels. + resolution, # Resolution of this block. + img_channels, # Number of input color channels. + first_layer_idx, # Index of the first layer. + architecture = 'resnet', # Architecture: 'orig', 'skip', 'resnet'. + activation = 'lrelu', # Activation function: 'relu', 'lrelu', etc. + resample_filter = [1,3,3,1], # Low-pass filter to apply when resampling activations. + conv_clamp = None, # Clamp the output of convolution layers to +-X, None = disable clamping. + use_fp16 = False, # Use FP16 for this block? + fp16_channels_last = False, # Use channels-last memory format with FP16? + freeze_layers = 0, # Freeze-D: Number of layers to freeze. + ): + assert in_channels in [0, tmp_channels] + assert architecture in ['orig', 'skip', 'resnet'] + super().__init__() + self.in_channels = in_channels + self.resolution = resolution + self.img_channels = img_channels + self.first_layer_idx = first_layer_idx + self.architecture = architecture + self.use_fp16 = use_fp16 + self.channels_last = (use_fp16 and fp16_channels_last) + self.register_buffer('resample_filter', upfirdn2d.setup_filter(resample_filter)) + + self.num_layers = 0 + def trainable_gen(): + while True: + layer_idx = self.first_layer_idx + self.num_layers + trainable = (layer_idx >= freeze_layers) + self.num_layers += 1 + yield trainable + trainable_iter = trainable_gen() + + if in_channels == 0 or architecture == 'skip': + self.fromrgb = Conv2dLayer(img_channels, tmp_channels, kernel_size=1, activation=activation, + trainable=next(trainable_iter), conv_clamp=conv_clamp, channels_last=self.channels_last) + + self.conv0 = Conv2dLayer(tmp_channels, tmp_channels, kernel_size=3, activation=activation, + trainable=next(trainable_iter), conv_clamp=conv_clamp, channels_last=self.channels_last) + + self.conv1 = Conv2dLayer(tmp_channels, out_channels, kernel_size=3, activation=activation, down=2, + trainable=next(trainable_iter), resample_filter=resample_filter, conv_clamp=conv_clamp, channels_last=self.channels_last) + + if architecture == 'resnet': + self.skip = Conv2dLayer(tmp_channels, out_channels, kernel_size=1, bias=False, down=2, + trainable=next(trainable_iter), resample_filter=resample_filter, channels_last=self.channels_last) + + def forward(self, x, img, force_fp32=False): + if (x if x is not None else img).device.type != 'cuda': + force_fp32 = True + dtype = torch.float16 if self.use_fp16 and not force_fp32 else torch.float32 + memory_format = torch.channels_last if self.channels_last and not force_fp32 else torch.contiguous_format + + # Input. + if x is not None: + misc.assert_shape(x, [None, self.in_channels, self.resolution, self.resolution]) + x = x.to(dtype=dtype, memory_format=memory_format) + + # FromRGB. + if self.in_channels == 0 or self.architecture == 'skip': + misc.assert_shape(img, [None, self.img_channels, self.resolution, self.resolution]) + img = img.to(dtype=dtype, memory_format=memory_format) + y = self.fromrgb(img) + x = x + y if x is not None else y + img = upfirdn2d.downsample2d(img, self.resample_filter) if self.architecture == 'skip' else None + + # Main layers. + if self.architecture == 'resnet': + y = self.skip(x, gain=np.sqrt(0.5)) + x = self.conv0(x) + x = self.conv1(x, gain=np.sqrt(0.5)) + x = y.add_(x) + else: + x = self.conv0(x) + x = self.conv1(x) + + assert x.dtype == dtype + return x, img + + def extra_repr(self): + return f'resolution={self.resolution:d}, architecture={self.architecture:s}' + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class MinibatchStdLayer(torch.nn.Module): + def __init__(self, group_size, num_channels=1): + super().__init__() + self.group_size = group_size + self.num_channels = num_channels + + def forward(self, x): + N, C, H, W = x.shape + with misc.suppress_tracer_warnings(): # as_tensor results are registered as constants + G = torch.min(torch.as_tensor(self.group_size), torch.as_tensor(N)) if self.group_size is not None else N + F = self.num_channels + c = C // F + + y = x.reshape(G, -1, F, c, H, W) # [GnFcHW] Split minibatch N into n groups of size G, and channels C into F groups of size c. + y = y - y.mean(dim=0) # [GnFcHW] Subtract mean over group. + y = y.square().mean(dim=0) # [nFcHW] Calc variance over group. + y = (y + 1e-8).sqrt() # [nFcHW] Calc stddev over group. + y = y.mean(dim=[2,3,4]) # [nF] Take average over channels and pixels. + y = y.reshape(-1, F, 1, 1) # [nF11] Add missing dimensions. + y = y.repeat(G, 1, H, W) # [NFHW] Replicate over group and pixels. + x = torch.cat([x, y], dim=1) # [NCHW] Append to input as new channels. + return x + + def extra_repr(self): + return f'group_size={self.group_size}, num_channels={self.num_channels:d}' + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class DiscriminatorEpilogue(torch.nn.Module): + def __init__(self, + in_channels, # Number of input channels. + cmap_dim, # Dimensionality of mapped conditioning label, 0 = no label. + resolution, # Resolution of this block. + img_channels, # Number of input color channels. + architecture = 'resnet', # Architecture: 'orig', 'skip', 'resnet'. + mbstd_group_size = 4, # Group size for the minibatch standard deviation layer, None = entire minibatch. + mbstd_num_channels = 1, # Number of features for the minibatch standard deviation layer, 0 = disable. + activation = 'lrelu', # Activation function: 'relu', 'lrelu', etc. + conv_clamp = None, # Clamp the output of convolution layers to +-X, None = disable clamping. + ): + assert architecture in ['orig', 'skip', 'resnet'] + super().__init__() + self.in_channels = in_channels + self.cmap_dim = cmap_dim + self.resolution = resolution + self.img_channels = img_channels + self.architecture = architecture + + if architecture == 'skip': + self.fromrgb = Conv2dLayer(img_channels, in_channels, kernel_size=1, activation=activation) + self.mbstd = MinibatchStdLayer(group_size=mbstd_group_size, num_channels=mbstd_num_channels) if mbstd_num_channels > 0 else None + self.conv = Conv2dLayer(in_channels + mbstd_num_channels, in_channels, kernel_size=3, activation=activation, conv_clamp=conv_clamp) + self.fc = FullyConnectedLayer(in_channels * (resolution ** 2), in_channels, activation=activation) + self.out = FullyConnectedLayer(in_channels, 1 if cmap_dim == 0 else cmap_dim) + + def forward(self, x, img, cmap, force_fp32=False): + misc.assert_shape(x, [None, self.in_channels, self.resolution, self.resolution]) # [NCHW] + _ = force_fp32 # unused + dtype = torch.float32 + memory_format = torch.contiguous_format + + # FromRGB. + x = x.to(dtype=dtype, memory_format=memory_format) + if self.architecture == 'skip': + misc.assert_shape(img, [None, self.img_channels, self.resolution, self.resolution]) + img = img.to(dtype=dtype, memory_format=memory_format) + x = x + self.fromrgb(img) + + # Main layers. + if self.mbstd is not None: + x = self.mbstd(x) + x = self.conv(x) + x = self.fc(x.flatten(1)) + x = self.out(x) + + # Conditioning. + if self.cmap_dim > 0: + misc.assert_shape(cmap, [None, self.cmap_dim]) + x = (x * cmap).sum(dim=1, keepdim=True) * (1 / np.sqrt(self.cmap_dim)) + + assert x.dtype == dtype + return x + + def extra_repr(self): + return f'resolution={self.resolution:d}, architecture={self.architecture:s}' + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class Discriminator(torch.nn.Module): + def __init__(self, + c_dim, # Conditioning label (C) dimensionality. + img_resolution, # Input resolution. + img_channels, # Number of input color channels. + architecture = 'resnet', # Architecture: 'orig', 'skip', 'resnet'. + channel_base = 32768, # Overall multiplier for the number of channels. + channel_max = 512, # Maximum number of channels in any layer. + num_fp16_res = 4, # Use FP16 for the N highest resolutions. + conv_clamp = 256, # Clamp the output of convolution layers to +-X, None = disable clamping. + cmap_dim = None, # Dimensionality of mapped conditioning label, None = default. + block_kwargs = {}, # Arguments for DiscriminatorBlock. + mapping_kwargs = {}, # Arguments for MappingNetwork. + epilogue_kwargs = {}, # Arguments for DiscriminatorEpilogue. + ): + super().__init__() + self.c_dim = c_dim + self.img_resolution = img_resolution + self.img_resolution_log2 = int(np.log2(img_resolution)) + self.img_channels = img_channels + self.block_resolutions = [2 ** i for i in range(self.img_resolution_log2, 2, -1)] + channels_dict = {res: min(channel_base // res, channel_max) for res in self.block_resolutions + [4]} + fp16_resolution = max(2 ** (self.img_resolution_log2 + 1 - num_fp16_res), 8) + + if cmap_dim is None: + cmap_dim = channels_dict[4] + if c_dim == 0: + cmap_dim = 0 + + common_kwargs = dict(img_channels=img_channels, architecture=architecture, conv_clamp=conv_clamp) + cur_layer_idx = 0 + for res in self.block_resolutions: + in_channels = channels_dict[res] if res < img_resolution else 0 + tmp_channels = channels_dict[res] + out_channels = channels_dict[res // 2] + use_fp16 = (res >= fp16_resolution) + block = DiscriminatorBlock(in_channels, tmp_channels, out_channels, resolution=res, + first_layer_idx=cur_layer_idx, use_fp16=use_fp16, **block_kwargs, **common_kwargs) + setattr(self, f'b{res}', block) + cur_layer_idx += block.num_layers + if c_dim > 0: + self.mapping = MappingNetwork(z_dim=0, c_dim=c_dim, w_dim=cmap_dim, num_ws=None, w_avg_beta=None, **mapping_kwargs) + self.b4 = DiscriminatorEpilogue(channels_dict[4], cmap_dim=cmap_dim, resolution=4, **epilogue_kwargs, **common_kwargs) + + def forward(self, img, c, update_emas=False, **block_kwargs): + _ = update_emas # unused + x = None + for res in self.block_resolutions: + block = getattr(self, f'b{res}') + x, img = block(x, img, **block_kwargs) + + cmap = None + if self.c_dim > 0: + cmap = self.mapping(None, c) + x = self.b4(x, img, cmap) + return x + + def extra_repr(self): + return f'c_dim={self.c_dim:d}, img_resolution={self.img_resolution:d}, img_channels={self.img_channels:d}' + +#---------------------------------------------------------------------------- \ No newline at end of file diff --git a/eg3d/training/networks_stylegan2_trip_and_vol.py b/eg3d/training/networks_stylegan2_trip_and_vol.py new file mode 100644 index 00000000..8e4dc489 --- /dev/null +++ b/eg3d/training/networks_stylegan2_trip_and_vol.py @@ -0,0 +1,1013 @@ +# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +"""Network architectures from the paper +"Analyzing and Improving the Image Quality of StyleGAN". +Matches the original implementation of configs E-F by Karras et al. at +https://github.com/NVlabs/stylegan2/blob/master/training/networks_stylegan2.py""" + +from ctypes.wintypes import PCHAR +from re import A +import numpy as np +import torch +from torch_utils import misc +from torch_utils import persistence +from torch_utils.ops import conv2d_resample +from torch_utils.ops import upfirdn2d +from torch_utils.ops import bias_act +from torch_utils.ops import fma + +from ipdb import set_trace as st +import multiprocessing +from torch_utils.utils_ds import grp_range_torch, parallel_FPS +import torch.nn as nn +import torch.nn.functional as F +import torch_scatter +import spconv.pytorch.conv as spconv +from training.costregnet import CostRegNet_Deeper +#---------------------------------------------------------------------------- + +@misc.profiled_function +def normalize_2nd_moment(x, dim=1, eps=1e-8): + return x * (x.square().mean(dim=dim, keepdim=True) + eps).rsqrt() + +#---------------------------------------------------------------------------- + +@misc.profiled_function +def modulated_conv2d( + x, # Input tensor of shape [batch_size, in_channels, in_height, in_width]. + weight, # Weight tensor of shape [out_channels, in_channels, kernel_height, kernel_width]. + styles, # Modulation coefficients of shape [batch_size, in_channels]. + noise = None, # Optional noise tensor to add to the output activations. + up = 1, # Integer upsampling factor. + down = 1, # Integer downsampling factor. + padding = 0, # Padding with respect to the upsampled image. + resample_filter = None, # Low-pass filter to apply when resampling activations. Must be prepared beforehand by calling upfirdn2d.setup_filter(). + demodulate = True, # Apply weight demodulation? + flip_weight = True, # False = convolution, True = correlation (matches torch.nn.functional.conv2d). + fused_modconv = True, # Perform modulation, convolution, and demodulation as a single fused operation? +): + batch_size = x.shape[0] + out_channels, in_channels, kh, kw = weight.shape + misc.assert_shape(weight, [out_channels, in_channels, kh, kw]) # [OIkk] + misc.assert_shape(x, [batch_size, in_channels, None, None]) # [NIHW] + misc.assert_shape(styles, [batch_size, in_channels]) # [NI] + + # Pre-normalize inputs to avoid FP16 overflow. + if x.dtype == torch.float16 and demodulate: + weight = weight * (1 / np.sqrt(in_channels * kh * kw) / weight.norm(float('inf'), dim=[1,2,3], keepdim=True)) # max_Ikk + styles = styles / styles.norm(float('inf'), dim=1, keepdim=True) # max_I + + # Calculate per-sample weights and demodulation coefficients. + w = None + dcoefs = None + if demodulate or fused_modconv: + w = weight.unsqueeze(0) # [NOIkk] + w = w * styles.reshape(batch_size, 1, -1, 1, 1) # [NOIkk] + if demodulate: + dcoefs = (w.square().sum(dim=[2,3,4]) + 1e-8).rsqrt() # [NO] + if demodulate and fused_modconv: + w = w * dcoefs.reshape(batch_size, -1, 1, 1, 1) # [NOIkk] + + # Execute by scaling the activations before and after the convolution. + if not fused_modconv: + x = x * styles.to(x.dtype).reshape(batch_size, -1, 1, 1) + x = conv2d_resample.conv2d_resample(x=x, w=weight.to(x.dtype), f=resample_filter, up=up, down=down, padding=padding, flip_weight=flip_weight) + if demodulate and noise is not None: + x = fma.fma(x, dcoefs.to(x.dtype).reshape(batch_size, -1, 1, 1), noise.to(x.dtype)) + elif demodulate: + x = x * dcoefs.to(x.dtype).reshape(batch_size, -1, 1, 1) + elif noise is not None: + x = x.add_(noise.to(x.dtype)) + return x + + # Execute as one fused op using grouped convolution. + with misc.suppress_tracer_warnings(): # this value will be treated as a constant + batch_size = int(batch_size) + misc.assert_shape(x, [batch_size, in_channels, None, None]) + x = x.reshape(1, -1, *x.shape[2:]) + w = w.reshape(-1, in_channels, kh, kw) + x = conv2d_resample.conv2d_resample(x=x, w=w.to(x.dtype), f=resample_filter, up=up, down=down, padding=padding, groups=batch_size, flip_weight=flip_weight) + x = x.reshape(batch_size, -1, *x.shape[2:]) + if noise is not None: + x = x.add_(noise) + return x + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class FullyConnectedLayer(torch.nn.Module): + def __init__(self, + in_features, # Number of input features. + out_features, # Number of output features. + bias = True, # Apply additive bias before the activation function? + activation = 'linear', # Activation function: 'relu', 'lrelu', etc. + lr_multiplier = 1, # Learning rate multiplier. + bias_init = 0, # Initial value for the additive bias. + ): + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.activation = activation + self.weight = torch.nn.Parameter(torch.randn([out_features, in_features]) / lr_multiplier) + self.bias = torch.nn.Parameter(torch.full([out_features], np.float32(bias_init))) if bias else None + self.weight_gain = lr_multiplier / np.sqrt(in_features) + self.bias_gain = lr_multiplier + + def forward(self, x): + + w = self.weight.to(x.dtype) * self.weight_gain + b = self.bias + if b is not None: + b = b.to(x.dtype) + if self.bias_gain != 1: + b = b * self.bias_gain + if self.activation == 'linear' and b is not None: + + x = torch.addmm(b.unsqueeze(0), x, w.t()) + else: + x = x.matmul(w.t()) + x = bias_act.bias_act(x, b, act=self.activation) + return x + + def extra_repr(self): + return f'in_features={self.in_features:d}, out_features={self.out_features:d}, activation={self.activation:s}' + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class Conv2dLayer(torch.nn.Module): + def __init__(self, + in_channels, # Number of input channels. + out_channels, # Number of output channels. + kernel_size, # Width and height of the convolution kernel. + bias = True, # Apply additive bias before the activation function? + activation = 'linear', # Activation function: 'relu', 'lrelu', etc. + up = 1, # Integer upsampling factor. + down = 1, # Integer downsampling factor. + resample_filter = [1,3,3,1], # Low-pass filter to apply when resampling activations. + conv_clamp = None, # Clamp the output to +-X, None = disable clamping. + channels_last = False, # Expect the input to have memory_format=channels_last? + trainable = True, # Update the weights of this layer during training? + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.activation = activation + self.up = up + self.down = down + self.conv_clamp = conv_clamp + self.register_buffer('resample_filter', upfirdn2d.setup_filter(resample_filter)) + self.padding = kernel_size // 2 + self.weight_gain = 1 / np.sqrt(in_channels * (kernel_size ** 2)) + self.act_gain = bias_act.activation_funcs[activation].def_gain + + memory_format = torch.channels_last if channels_last else torch.contiguous_format + weight = torch.randn([out_channels, in_channels, kernel_size, kernel_size]).to(memory_format=memory_format) + bias = torch.zeros([out_channels]) if bias else None + if trainable: + self.weight = torch.nn.Parameter(weight) + self.bias = torch.nn.Parameter(bias) if bias is not None else None + else: + self.register_buffer('weight', weight) + if bias is not None: + self.register_buffer('bias', bias) + else: + self.bias = None + + def forward(self, x, gain=1): + w = self.weight * self.weight_gain + b = self.bias.to(x.dtype) if self.bias is not None else None + flip_weight = (self.up == 1) # slightly faster + x = conv2d_resample.conv2d_resample(x=x, w=w.to(x.dtype), f=self.resample_filter, up=self.up, down=self.down, padding=self.padding, flip_weight=flip_weight) + + act_gain = self.act_gain * gain + act_clamp = self.conv_clamp * gain if self.conv_clamp is not None else None + x = bias_act.bias_act(x, b, act=self.activation, gain=act_gain, clamp=act_clamp) + return x + + def extra_repr(self): + return ' '.join([ + f'in_channels={self.in_channels:d}, out_channels={self.out_channels:d}, activation={self.activation:s},', + f'up={self.up}, down={self.down}']) + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class MappingNetwork(torch.nn.Module): + def __init__(self, + z_dim, # Input latent (Z) dimensionality, 0 = no latent. + c_dim, # Conditioning label (C) dimensionality, 0 = no label. + w_dim, # Intermediate latent (W) dimensionality. + num_ws, # Number of intermediate latents to output, None = do not broadcast. + num_layers = 8, # Number of mapping layers. + embed_features = None, # Label embedding dimensionality, None = same as w_dim. + layer_features = None, # Number of intermediate features in the mapping layers, None = same as w_dim. + activation = 'lrelu', # Activation function: 'relu', 'lrelu', etc. + lr_multiplier = 0.01, # Learning rate multiplier for the mapping layers. + w_avg_beta = 0.998, # Decay for tracking the moving average of W during training, None = do not track. + ): + super().__init__() + self.z_dim = z_dim + self.c_dim = c_dim + self.w_dim = w_dim + self.num_ws = num_ws + self.num_layers = num_layers + self.w_avg_beta = w_avg_beta + + if embed_features is None: + embed_features = w_dim + if c_dim == 0: + embed_features = 0 + if layer_features is None: + layer_features = w_dim + features_list = [z_dim + embed_features] + [layer_features] * (num_layers - 1) + [w_dim] + + if c_dim > 0: + self.embed = FullyConnectedLayer(c_dim, embed_features) + for idx in range(num_layers): + in_features = features_list[idx] + out_features = features_list[idx + 1] + layer = FullyConnectedLayer(in_features, out_features, activation=activation, lr_multiplier=lr_multiplier) + setattr(self, f'fc{idx}', layer) + + if num_ws is not None and w_avg_beta is not None: + self.register_buffer('w_avg', torch.zeros([w_dim])) + + def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, update_emas=False): + # Embed, normalize, and concat inputs. + x = None + with torch.autograd.profiler.record_function('input'): + if self.z_dim > 0: + misc.assert_shape(z, [None, self.z_dim]) + x = normalize_2nd_moment(z.to(torch.float32)) + if self.c_dim > 0: + misc.assert_shape(c, [None, self.c_dim]) + y = normalize_2nd_moment(self.embed(c.to(torch.float32))) + x = torch.cat([x, y], dim=1) if x is not None else y + + # Main layers. + for idx in range(self.num_layers): + layer = getattr(self, f'fc{idx}') + x = layer(x) + + # Update moving average of W. + if update_emas and self.w_avg_beta is not None: + with torch.autograd.profiler.record_function('update_w_avg'): + self.w_avg.copy_(x.detach().mean(dim=0).lerp(self.w_avg, self.w_avg_beta)) + + # Broadcast. + if self.num_ws is not None: + with torch.autograd.profiler.record_function('broadcast'): + x = x.unsqueeze(1).repeat([1, self.num_ws, 1]) + + # Apply truncation. + if truncation_psi != 1: + with torch.autograd.profiler.record_function('truncate'): + assert self.w_avg_beta is not None + if self.num_ws is None or truncation_cutoff is None: + x = self.w_avg.lerp(x, truncation_psi) + else: + x[:, :truncation_cutoff] = self.w_avg.lerp(x[:, :truncation_cutoff], truncation_psi) + return x + + def extra_repr(self): + return f'z_dim={self.z_dim:d}, c_dim={self.c_dim:d}, w_dim={self.w_dim:d}, num_ws={self.num_ws:d}' + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class SynthesisLayer(torch.nn.Module): + def __init__(self, + in_channels, # Number of input channels. + out_channels, # Number of output channels. + w_dim, # Intermediate latent (W) dimensionality. + resolution, # Resolution of this layer. + kernel_size = 3, # Convolution kernel size. + up = 1, # Integer upsampling factor. + use_noise = True, # Enable noise input? + activation = 'lrelu', # Activation function: 'relu', 'lrelu', etc. + resample_filter = [1,3,3,1], # Low-pass filter to apply when resampling activations. + conv_clamp = None, # Clamp the output of convolution layers to +-X, None = disable clamping. + channels_last = False, # Use channels_last format for the weights? + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.w_dim = w_dim + self.resolution = resolution + self.up = up + self.use_noise = use_noise + self.activation = activation + self.conv_clamp = conv_clamp + self.register_buffer('resample_filter', upfirdn2d.setup_filter(resample_filter)) + self.padding = kernel_size // 2 + self.act_gain = bias_act.activation_funcs[activation].def_gain + + self.affine = FullyConnectedLayer(w_dim, in_channels, bias_init=1) + memory_format = torch.channels_last if channels_last else torch.contiguous_format + self.weight = torch.nn.Parameter(torch.randn([out_channels, in_channels, kernel_size, kernel_size]).to(memory_format=memory_format)) + if use_noise: + self.register_buffer('noise_const', torch.randn([resolution, resolution])) + self.noise_strength = torch.nn.Parameter(torch.zeros([])) + self.bias = torch.nn.Parameter(torch.zeros([out_channels])) + + def forward(self, x, w, noise_mode='random', fused_modconv=True, gain=1): + assert noise_mode in ['random', 'const', 'none'] + in_resolution = self.resolution // self.up + misc.assert_shape(x, [None, self.in_channels, in_resolution, in_resolution]) + styles = self.affine(w) + + noise = None + if self.use_noise and noise_mode == 'random': + noise = torch.randn([x.shape[0], 1, self.resolution, self.resolution], device=x.device) * self.noise_strength + if self.use_noise and noise_mode == 'const': + noise = self.noise_const * self.noise_strength + + flip_weight = (self.up == 1) # slightly faster + x = modulated_conv2d(x=x, weight=self.weight, styles=styles, noise=noise, up=self.up, + padding=self.padding, resample_filter=self.resample_filter, flip_weight=flip_weight, fused_modconv=fused_modconv) + + act_gain = self.act_gain * gain + act_clamp = self.conv_clamp * gain if self.conv_clamp is not None else None + x = bias_act.bias_act(x, self.bias.to(x.dtype), act=self.activation, gain=act_gain, clamp=act_clamp) + return x + + def extra_repr(self): + return ' '.join([ + f'in_channels={self.in_channels:d}, out_channels={self.out_channels:d}, w_dim={self.w_dim:d},', + f'resolution={self.resolution:d}, up={self.up}, activation={self.activation:s}']) + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class ToRGBLayer(torch.nn.Module): + def __init__(self, in_channels, out_channels, w_dim, kernel_size=1, conv_clamp=None, channels_last=False): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.w_dim = w_dim + self.conv_clamp = conv_clamp + self.affine = FullyConnectedLayer(w_dim, in_channels, bias_init=1) + memory_format = torch.channels_last if channels_last else torch.contiguous_format + self.weight = torch.nn.Parameter(torch.randn([out_channels, in_channels, kernel_size, kernel_size]).to(memory_format=memory_format)) + self.bias = torch.nn.Parameter(torch.zeros([out_channels])) + self.weight_gain = 1 / np.sqrt(in_channels * (kernel_size ** 2)) + + def forward(self, x, w, fused_modconv=True): + styles = self.affine(w) * self.weight_gain + x = modulated_conv2d(x=x, weight=self.weight, styles=styles, demodulate=False, fused_modconv=fused_modconv) + x = bias_act.bias_act(x, self.bias.to(x.dtype), clamp=self.conv_clamp) + return x + + def extra_repr(self): + return f'in_channels={self.in_channels:d}, out_channels={self.out_channels:d}, w_dim={self.w_dim:d}' + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class SynthesisBlock(torch.nn.Module): + def __init__(self, + in_channels, # Number of input channels, 0 = first block. + out_channels, # Number of output channels. + w_dim, # Intermediate latent (W) dimensionality. + resolution, # Resolution of this block. + img_channels, # Number of output color channels. + is_last, # Is this the last block? + architecture = 'skip', # Architecture: 'orig', 'skip', 'resnet'. + resample_filter = [1,3,3,1], # Low-pass filter to apply when resampling activations. + conv_clamp = 256, # Clamp the output of convolution layers to +-X, None = disable clamping. + use_fp16 = False, # Use FP16 for this block? + fp16_channels_last = False, # Use channels-last memory format with FP16? + fused_modconv_default = True, # Default value of fused_modconv. 'inference_only' = True for inference, False for training. + **layer_kwargs, # Arguments for SynthesisLayer. + ): + assert architecture in ['orig', 'skip', 'resnet'] + super().__init__() + self.in_channels = in_channels + self.w_dim = w_dim + self.resolution = resolution + self.img_channels = img_channels + self.is_last = is_last + self.architecture = architecture + self.use_fp16 = use_fp16 + self.channels_last = (use_fp16 and fp16_channels_last) + self.fused_modconv_default = fused_modconv_default + self.register_buffer('resample_filter', upfirdn2d.setup_filter(resample_filter)) + self.num_conv = 0 + self.num_torgb = 0 + + if in_channels == 0: + self.const = torch.nn.Parameter(torch.randn([out_channels, resolution, resolution])) + + if in_channels != 0: + self.conv0 = SynthesisLayer(in_channels, out_channels, w_dim=w_dim, resolution=resolution, up=2, + resample_filter=resample_filter, conv_clamp=conv_clamp, channels_last=self.channels_last, **layer_kwargs) + self.num_conv += 1 + + self.conv1 = SynthesisLayer(out_channels, out_channels, w_dim=w_dim, resolution=resolution, + conv_clamp=conv_clamp, channels_last=self.channels_last, **layer_kwargs) + self.num_conv += 1 + + if is_last or architecture == 'skip': + self.torgb = ToRGBLayer(out_channels, img_channels, w_dim=w_dim, + conv_clamp=conv_clamp, channels_last=self.channels_last) + self.num_torgb += 1 + + if in_channels != 0 and architecture == 'resnet': + self.skip = Conv2dLayer(in_channels, out_channels, kernel_size=1, bias=False, up=2, + resample_filter=resample_filter, channels_last=self.channels_last) + + def forward(self, x, img, ws, force_fp32=False, fused_modconv=None, update_emas=False, **layer_kwargs): + _ = update_emas # unused + misc.assert_shape(ws, [None, self.num_conv + self.num_torgb, self.w_dim]) + w_iter = iter(ws.unbind(dim=1)) + if ws.device.type != 'cuda': + force_fp32 = True + dtype = torch.float16 if self.use_fp16 and not force_fp32 else torch.float32 + memory_format = torch.channels_last if self.channels_last and not force_fp32 else torch.contiguous_format + if fused_modconv is None: + fused_modconv = self.fused_modconv_default + if fused_modconv == 'inference_only': + fused_modconv = (not self.training) + + # Input. + if self.in_channels == 0: + x = self.const.to(dtype=dtype, memory_format=memory_format) + x = x.unsqueeze(0).repeat([ws.shape[0], 1, 1, 1]) + else: + misc.assert_shape(x, [None, self.in_channels, self.resolution // 2, self.resolution // 2]) + x = x.to(dtype=dtype, memory_format=memory_format) + + # Main layers. + if self.in_channels == 0: + x = self.conv1(x, next(w_iter), fused_modconv=fused_modconv, **layer_kwargs) + elif self.architecture == 'resnet': + y = self.skip(x, gain=np.sqrt(0.5)) + x = self.conv0(x, next(w_iter), fused_modconv=fused_modconv, **layer_kwargs) + x = self.conv1(x, next(w_iter), fused_modconv=fused_modconv, gain=np.sqrt(0.5), **layer_kwargs) + x = y.add_(x) + else: + x = self.conv0(x, next(w_iter), fused_modconv=fused_modconv, **layer_kwargs) + x = self.conv1(x, next(w_iter), fused_modconv=fused_modconv, **layer_kwargs) + + # ToRGB. + if img is not None: + misc.assert_shape(img, [None, self.img_channels, self.resolution // 2, self.resolution // 2]) + img = upfirdn2d.upsample2d(img, self.resample_filter) + if self.is_last or self.architecture == 'skip': + y = self.torgb(x, next(w_iter), fused_modconv=fused_modconv) + y = y.to(dtype=torch.float32, memory_format=torch.contiguous_format) + img = img.add_(y) if img is not None else y + + assert x.dtype == dtype + assert img is None or img.dtype == torch.float32 + return x, img + + def extra_repr(self): + return f'resolution={self.resolution:d}, architecture={self.architecture:s}' + +#---------------------------------------------------------------------------- + +# directly modify this to generate volume +@persistence.persistent_class +class SynthesisNetwork(torch.nn.Module): + def __init__(self, + w_dim, # Intermediate latent (W) dimensionality. + img_resolution, # Output image resolution. + img_channels, # Number of color channels. + channel_base = 32768, # Overall multiplier for the number of channels. + channel_max = 512, # Maximum number of channels in any layer. + num_fp16_res = 4, # Use FP16 for the N highest resolutions. + **block_kwargs, # Arguments for SynthesisBlock. + ): + assert img_resolution >= 4 and img_resolution & (img_resolution - 1) == 0 + super().__init__() + self.w_dim = w_dim + self.img_resolution = img_resolution + self.img_resolution_log2 = int(np.log2(img_resolution)) + self.img_channels = img_channels + self.num_fp16_res = num_fp16_res + self.block_resolutions = [2 ** i for i in range(2, self.img_resolution_log2 + 1)] + channels_dict = {res: min(channel_base // res, channel_max) for res in self.block_resolutions} + fp16_resolution = max(2 ** (self.img_resolution_log2 + 1 - num_fp16_res), 8) + + self.num_ws = 0 + for res in self.block_resolutions: + in_channels = channels_dict[res // 2] if res > 4 else 0 + out_channels = channels_dict[res] + use_fp16 = (res >= fp16_resolution) + is_last = (res == self.img_resolution) + block = SynthesisBlock(in_channels, out_channels, w_dim=w_dim, resolution=res, + img_channels=img_channels, is_last=is_last, use_fp16=use_fp16, **block_kwargs) + self.num_ws += block.num_conv + if is_last: + self.num_ws += block.num_torgb + setattr(self, f'b{res}', block) + + ###### hard-code attr for voxelize ####### + self.vfe_model = PointNet(fea_dim=9, out_pt_fea_dim=32) # TODO: modify this hard-coded thing + self.pt_selection = 'random' + self.max_pt = 256 + self.pt_pooling = 'max' + if self.pt_pooling == 'max': + self.pool_dim = 64 + self.pos_enc_dim = 0 + # self.fea_compre = 32 + # self.fea_compression = nn.Sequential( + # nn.Linear(self.pool_dim+self.pos_enc_dim, self.fea_compre), + # nn.ReLU() + # ).cuda() + ######### for unet3d ############ + unet_in_channels = 32 + self.unet3d=CostRegNet_Deeper(unet_in_channels, norm_act= nn.BatchNorm3d).to(torch.device("cuda")) + _grid_size=64 + self.grid_size = [_grid_size]*3 + + def forward(self, ws, pc, box_warp, **block_kwargs): + # def forward(self, ws, **block_kwargs): + RETURN_IMG=True + RETURN_BOTH=True + if RETURN_IMG or RETURN_BOTH: + + block_ws = [] + ######## latents ---------------- + + with torch.autograd.profiler.record_function('split_ws'): + misc.assert_shape(ws, [None, self.num_ws, self.w_dim]) + ws = ws.to(torch.float32) + w_idx = 0 + for res in self.block_resolutions: + block = getattr(self, f'b{res}') + block_ws.append(ws.narrow(1, w_idx, block.num_conv + block.num_torgb)) + w_idx += block.num_conv + + + # ----change to all with the same global latent------------ + + # v1: no need to process ws: concat all at the bottleneck + # v2: progressively add latents during upconv + + + ########## generate tri-plane ############## + + x = img = None + # st() # pc.shape + for res, cur_ws in zip(self.block_resolutions, block_ws): + block = getattr(self, f'b{res}') + x, img = block(x, img, cur_ws, **block_kwargs) + # st() # align with img.shape: torch.Size([4, 96, 256, 256]): B,C,H,W + if not RETURN_BOTH: + return img + + # ----change to 3D Unet ------------ + + # # 1. voxelize input pc + + B,_,_=pc.shape + _coor, _feature_3d, density_volume, voxel_size = self.voxelize_spconv_sparse_pointnet( + pc=pc, box_warp=box_warp,grid_size=self.grid_size) + _ret = spconv.SparseConvTensor(_feature_3d, _coor.int(), np.array(self.grid_size), + B) # sp_tensor batch = B*V + _feature_3d = _ret.dense(channels_first = True).contiguous() # [B, C, X, Y, Z], C=32 + + # voxelize_spconv_sparse_pointnet(self, batch_pcl, grid_size=[], + # batch_bbox=None, pointnet_input=None, pyramid_layer=None): + # # 2. 3D Unet: special: with addtional input ws to concate at the bottle neck so that the upconv part can serve as generator + # self.backbone as in mvsnerf.models + _feature_3d = self.unet3d(_feature_3d.contiguous()) # 3d CONV takes [B, C, X, Y, Z] as input + # st() + volume = _feature_3d.permute(0,1,4,3,2) + + if not RETURN_IMG: + return volume + else: + assert RETURN_BOTH + # st() + return img, volume + + def extra_repr(self): + return ' '.join([ + f'w_dim={self.w_dim:d}, num_ws={self.num_ws:d},', + f'img_resolution={self.img_resolution:d}, img_channels={self.img_channels:d},', + f'num_fp16_res={self.num_fp16_res:d}']) + + + def voxelize_spconv_sparse_pointnet(self, pc, grid_size=[], + box_warp=None, pointnet_input=None): + ######## parameter alignment ######### + batch_pcl = pc[...,:3].unsqueeze(1) + batch_mtl = pc[...,3:].unsqueeze(1) + device=batch_pcl.device + B,V,_,_ = batch_pcl.shape # torch.Size([4, 1, 1500, 9]) + #check box_warp.shape + batch_bbox = torch.tensor([ + [-box_warp/2, -box_warp/2, -box_warp/2], + [box_warp/2, box_warp/2, box_warp/2] + ],device=batch_pcl.device)[None, None,...].repeat(B,V,1,1) + + pointnet_input='local_xyz' # hard-code + feature ='pointnet' + + ######## function logic ######### + + grid_size = torch.tensor(grid_size, device=device) + + ## direct batch voxelization + + batch_voxel_size = (batch_bbox[:,:,1:]-batch_bbox[:,:,:1])/grid_size + voxel_size = batch_voxel_size + batch_xyz_cube_pos = torch.div((batch_pcl-batch_bbox[:,:,:1]), batch_voxel_size, rounding_mode='floor') + if pointnet_input== 'local_xyz': + batch_pcl_local = (batch_pcl - batch_bbox[:,:,:1] - batch_xyz_cube_pos*batch_voxel_size) / batch_voxel_size - 0.5 + batch_pcl_local = torch.cat([batch_pcl_local, batch_mtl], dim=-1) #torch.Size([4, 1, 1500, 9]) + + cat_pt_fea, cat_pt_ind = [], [] + for i_batch in range(len(batch_xyz_cube_pos)): + for i_view in range(len(batch_xyz_cube_pos[i_batch])): + cat_pt_fea.append(batch_pcl_local[i_batch, i_view]) + cat_pt_ind.append(F.pad(batch_xyz_cube_pos[i_batch, i_view],(1,0),'constant',value = i_batch*V+i_view)) + cat_pt_fea = torch.cat(cat_pt_fea,dim=0) + cat_pt_ind = torch.cat(cat_pt_ind,dim = 0) + # st() + else: + raise NotImplemented(False) + + pt_num = cat_pt_ind.shape[0] + # shuffle the data + cur_dev = cat_pt_fea.get_device() + shuffled_ind = torch.randperm(pt_num,device = cur_dev) + cat_pt_fea = cat_pt_fea[shuffled_ind,:] + cat_pt_ind = cat_pt_ind[shuffled_ind,:] + + # unique xy grid index + unq, unq_inv, unq_cnt = torch.unique(cat_pt_ind,return_inverse=True, return_counts=True, dim=0) + unq = unq.type(torch.int64) + + + # subsample pts + if self.pt_selection == 'random': + grp_ind = grp_range_torch(unq_cnt,cur_dev)[torch.argsort(torch.argsort(unq_inv))] # convert the array that is in the order of grid to the order of cat_pt_feature + remain_ind = grp_ind < self.max_pt # randomly sample max_pt points inside a grid + elif self.pt_selection == 'farthest': + unq_ind = np.split(np.argsort(unq_inv.detach().cpu().numpy()), np.cumsum(unq_cnt.detach().cpu().numpy()[:-1])) + remain_ind = np.zeros((pt_num,),dtype = np.bool) + np_cat_fea = cat_pt_fea.detach().cpu().numpy()[:,:3] + pool_in = [] + for i_inds in unq_ind: + if len(i_inds) > self.max_pt: + pool_in.append((np_cat_fea[i_inds,:],self.max_pt)) + if len(pool_in) > 0: + pool = multiprocessing.Pool(multiprocessing.cpu_count()) + FPS_results = pool.starmap(parallel_FPS, pool_in) + pool.close() + pool.join() + count = 0 + for i_inds in unq_ind: + if len(i_inds) <= self.max_pt: + remain_ind[i_inds] = True + else: + remain_ind[i_inds[FPS_results[count]]] = True + count += 1 + + cat_pt_fea = cat_pt_fea[remain_ind,:] + cat_pt_ind = cat_pt_ind[remain_ind,:] + unq_inv = unq_inv[remain_ind] + unq_cnt = torch.clamp(unq_cnt,max=self.max_pt) + # construct density volume from unqcnt + batch_densities_volumes = unq_cnt[...,None] # add one more dim + + if feature=='embedding': + processed_cat_pt_fea = self.voxel_embed(cat_pt_fea) + elif feature=='pointnet': + processed_cat_pt_fea = self.vfe_model(cat_pt_fea) # global pointnet + # st() + + if self.pt_pooling == 'max': + pooled_data = torch_scatter.scatter_max(processed_cat_pt_fea, unq_inv, dim=0)[0] # choose the max feature for each grid + else: raise NotImplementedError + + # if self.fea_compre: + # processed_pooled_data = self.fea_compression(pooled_data) + # else: + processed_pooled_data = pooled_data + + return unq, processed_pooled_data, batch_densities_volumes, voxel_size + + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class PointNet(torch.nn.Module): + # def __init__(self, cfg): + def __init__(self, fea_dim, out_pt_fea_dim): + super().__init__() + # fea_dim = cfg.DATA_CONFIG.DATALOADER.DATA_DIM + # out_pt_fea_dim = cfg.MODEL.VFE.OUT_CHANNEL + + self.PPmodel = nn.Sequential( + # nn.BatchNorm1d(fea_dim), + nn.Linear(fea_dim, 64), + nn.BatchNorm1d(64), + nn.ReLU(), + nn.Linear(64, 128), + nn.BatchNorm1d(128), + nn.ReLU(), + nn.Linear(128, 256), + nn.BatchNorm1d(256), + nn.ReLU(), + nn.Linear(256, out_pt_fea_dim) + ) + + def forward(self, x): + return self.PPmodel(x) + +#---------------------------------------------------------------------------- + + + +@persistence.persistent_class +class Generator(torch.nn.Module): + def __init__(self, + z_dim, # Input latent (Z) dimensionality. + c_dim, # Conditioning label (C) dimensionality. + w_dim, # Intermediate latent (W) dimensionality. + ####### newly added parameters ###### + pc_dim, # Conditioning poincloud (PC) dimensionality. + volume_res, # Volume resolution. + ########################################## + img_resolution, # Output resolution. + img_channels, # Number of output color channels. + mapping_kwargs = {}, # Arguments for MappingNetwork. + **synthesis_kwargs, # Arguments for SynthesisNetwork. + ): + super().__init__() + self.z_dim = z_dim + self.c_dim = c_dim + self.w_dim = w_dim + ####### newly added parameters ###### + self.pc_dim=pc_dim + self.volume_res=volume_res + ########################################## + self.img_resolution = img_resolution + self.img_channels = img_channels + self.synthesis = SynthesisNetwork(w_dim=w_dim, img_resolution=img_resolution, img_channels=img_channels, **synthesis_kwargs) + self.num_ws = self.synthesis.num_ws + self.mapping = MappingNetwork(z_dim=z_dim, c_dim=c_dim, w_dim=w_dim, num_ws=self.num_ws, **mapping_kwargs) + + def forward(self, z, c, pc, truncation_psi=1, truncation_cutoff=None, update_emas=False, **synthesis_kwargs): + # TODO: whether to include pc info during self.mapping?? + ws = self.mapping(z, c, truncation_psi=truncation_psi, truncation_cutoff=truncation_cutoff, update_emas=update_emas) + img = self.synthesis(ws, pc, update_emas=update_emas, **synthesis_kwargs) + return img + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class DiscriminatorBlock(torch.nn.Module): + def __init__(self, + in_channels, # Number of input channels, 0 = first block. + tmp_channels, # Number of intermediate channels. + out_channels, # Number of output channels. + resolution, # Resolution of this block. + img_channels, # Number of input color channels. + first_layer_idx, # Index of the first layer. + architecture = 'resnet', # Architecture: 'orig', 'skip', 'resnet'. + activation = 'lrelu', # Activation function: 'relu', 'lrelu', etc. + resample_filter = [1,3,3,1], # Low-pass filter to apply when resampling activations. + conv_clamp = None, # Clamp the output of convolution layers to +-X, None = disable clamping. + use_fp16 = False, # Use FP16 for this block? + fp16_channels_last = False, # Use channels-last memory format with FP16? + freeze_layers = 0, # Freeze-D: Number of layers to freeze. + ): + assert in_channels in [0, tmp_channels] + assert architecture in ['orig', 'skip', 'resnet'] + super().__init__() + self.in_channels = in_channels + self.resolution = resolution + self.img_channels = img_channels + self.first_layer_idx = first_layer_idx + self.architecture = architecture + self.use_fp16 = use_fp16 + self.channels_last = (use_fp16 and fp16_channels_last) + self.register_buffer('resample_filter', upfirdn2d.setup_filter(resample_filter)) + + self.num_layers = 0 + def trainable_gen(): + while True: + layer_idx = self.first_layer_idx + self.num_layers + trainable = (layer_idx >= freeze_layers) + self.num_layers += 1 + yield trainable + trainable_iter = trainable_gen() + + if in_channels == 0 or architecture == 'skip': + self.fromrgb = Conv2dLayer(img_channels, tmp_channels, kernel_size=1, activation=activation, + trainable=next(trainable_iter), conv_clamp=conv_clamp, channels_last=self.channels_last) + + self.conv0 = Conv2dLayer(tmp_channels, tmp_channels, kernel_size=3, activation=activation, + trainable=next(trainable_iter), conv_clamp=conv_clamp, channels_last=self.channels_last) + + self.conv1 = Conv2dLayer(tmp_channels, out_channels, kernel_size=3, activation=activation, down=2, + trainable=next(trainable_iter), resample_filter=resample_filter, conv_clamp=conv_clamp, channels_last=self.channels_last) + + if architecture == 'resnet': + self.skip = Conv2dLayer(tmp_channels, out_channels, kernel_size=1, bias=False, down=2, + trainable=next(trainable_iter), resample_filter=resample_filter, channels_last=self.channels_last) + + def forward(self, x, img, force_fp32=False): + if (x if x is not None else img).device.type != 'cuda': + force_fp32 = True + dtype = torch.float16 if self.use_fp16 and not force_fp32 else torch.float32 + memory_format = torch.channels_last if self.channels_last and not force_fp32 else torch.contiguous_format + + # Input. + if x is not None: + misc.assert_shape(x, [None, self.in_channels, self.resolution, self.resolution]) + x = x.to(dtype=dtype, memory_format=memory_format) + + # FromRGB. + if self.in_channels == 0 or self.architecture == 'skip': + misc.assert_shape(img, [None, self.img_channels, self.resolution, self.resolution]) + img = img.to(dtype=dtype, memory_format=memory_format) + y = self.fromrgb(img) + x = x + y if x is not None else y + img = upfirdn2d.downsample2d(img, self.resample_filter) if self.architecture == 'skip' else None + + # Main layers. + if self.architecture == 'resnet': + y = self.skip(x, gain=np.sqrt(0.5)) + x = self.conv0(x) + x = self.conv1(x, gain=np.sqrt(0.5)) + x = y.add_(x) + else: + x = self.conv0(x) + x = self.conv1(x) + + assert x.dtype == dtype + return x, img + + def extra_repr(self): + return f'resolution={self.resolution:d}, architecture={self.architecture:s}' + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class MinibatchStdLayer(torch.nn.Module): + def __init__(self, group_size, num_channels=1): + super().__init__() + self.group_size = group_size + self.num_channels = num_channels + + def forward(self, x): + N, C, H, W = x.shape + with misc.suppress_tracer_warnings(): # as_tensor results are registered as constants + G = torch.min(torch.as_tensor(self.group_size), torch.as_tensor(N)) if self.group_size is not None else N + F = self.num_channels + c = C // F + + y = x.reshape(G, -1, F, c, H, W) # [GnFcHW] Split minibatch N into n groups of size G, and channels C into F groups of size c. + y = y - y.mean(dim=0) # [GnFcHW] Subtract mean over group. + y = y.square().mean(dim=0) # [nFcHW] Calc variance over group. + y = (y + 1e-8).sqrt() # [nFcHW] Calc stddev over group. + y = y.mean(dim=[2,3,4]) # [nF] Take average over channels and pixels. + y = y.reshape(-1, F, 1, 1) # [nF11] Add missing dimensions. + y = y.repeat(G, 1, H, W) # [NFHW] Replicate over group and pixels. + x = torch.cat([x, y], dim=1) # [NCHW] Append to input as new channels. + return x + + def extra_repr(self): + return f'group_size={self.group_size}, num_channels={self.num_channels:d}' + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class DiscriminatorEpilogue(torch.nn.Module): + def __init__(self, + in_channels, # Number of input channels. + cmap_dim, # Dimensionality of mapped conditioning label, 0 = no label. + resolution, # Resolution of this block. + img_channels, # Number of input color channels. + architecture = 'resnet', # Architecture: 'orig', 'skip', 'resnet'. + mbstd_group_size = 4, # Group size for the minibatch standard deviation layer, None = entire minibatch. + mbstd_num_channels = 1, # Number of features for the minibatch standard deviation layer, 0 = disable. + activation = 'lrelu', # Activation function: 'relu', 'lrelu', etc. + conv_clamp = None, # Clamp the output of convolution layers to +-X, None = disable clamping. + ): + assert architecture in ['orig', 'skip', 'resnet'] + super().__init__() + self.in_channels = in_channels + self.cmap_dim = cmap_dim + self.resolution = resolution + self.img_channels = img_channels + self.architecture = architecture + + if architecture == 'skip': + self.fromrgb = Conv2dLayer(img_channels, in_channels, kernel_size=1, activation=activation) + self.mbstd = MinibatchStdLayer(group_size=mbstd_group_size, num_channels=mbstd_num_channels) if mbstd_num_channels > 0 else None + self.conv = Conv2dLayer(in_channels + mbstd_num_channels, in_channels, kernel_size=3, activation=activation, conv_clamp=conv_clamp) + self.fc = FullyConnectedLayer(in_channels * (resolution ** 2), in_channels, activation=activation) + self.out = FullyConnectedLayer(in_channels, 1 if cmap_dim == 0 else cmap_dim) + + def forward(self, x, img, cmap, force_fp32=False): + misc.assert_shape(x, [None, self.in_channels, self.resolution, self.resolution]) # [NCHW] + _ = force_fp32 # unused + dtype = torch.float32 + memory_format = torch.contiguous_format + + # FromRGB. + x = x.to(dtype=dtype, memory_format=memory_format) + if self.architecture == 'skip': + misc.assert_shape(img, [None, self.img_channels, self.resolution, self.resolution]) + img = img.to(dtype=dtype, memory_format=memory_format) + x = x + self.fromrgb(img) + + # Main layers. + if self.mbstd is not None: + x = self.mbstd(x) + x = self.conv(x) + x = self.fc(x.flatten(1)) + x = self.out(x) + + # Conditioning. + if self.cmap_dim > 0: + misc.assert_shape(cmap, [None, self.cmap_dim]) + x = (x * cmap).sum(dim=1, keepdim=True) * (1 / np.sqrt(self.cmap_dim)) + + assert x.dtype == dtype + return x + + def extra_repr(self): + return f'resolution={self.resolution:d}, architecture={self.architecture:s}' + +#---------------------------------------------------------------------------- + +@persistence.persistent_class +class Discriminator(torch.nn.Module): + def __init__(self, + c_dim, # Conditioning label (C) dimensionality. + img_resolution, # Input resolution. + img_channels, # Number of input color channels. + architecture = 'resnet', # Architecture: 'orig', 'skip', 'resnet'. + channel_base = 32768, # Overall multiplier for the number of channels. + channel_max = 512, # Maximum number of channels in any layer. + num_fp16_res = 4, # Use FP16 for the N highest resolutions. + conv_clamp = 256, # Clamp the output of convolution layers to +-X, None = disable clamping. + cmap_dim = None, # Dimensionality of mapped conditioning label, None = default. + block_kwargs = {}, # Arguments for DiscriminatorBlock. + mapping_kwargs = {}, # Arguments for MappingNetwork. + epilogue_kwargs = {}, # Arguments for DiscriminatorEpilogue. + ): + super().__init__() + self.c_dim = c_dim + self.img_resolution = img_resolution + self.img_resolution_log2 = int(np.log2(img_resolution)) + self.img_channels = img_channels + self.block_resolutions = [2 ** i for i in range(self.img_resolution_log2, 2, -1)] + channels_dict = {res: min(channel_base // res, channel_max) for res in self.block_resolutions + [4]} + fp16_resolution = max(2 ** (self.img_resolution_log2 + 1 - num_fp16_res), 8) + + if cmap_dim is None: + cmap_dim = channels_dict[4] + if c_dim == 0: + cmap_dim = 0 + + common_kwargs = dict(img_channels=img_channels, architecture=architecture, conv_clamp=conv_clamp) + cur_layer_idx = 0 + for res in self.block_resolutions: + in_channels = channels_dict[res] if res < img_resolution else 0 + tmp_channels = channels_dict[res] + out_channels = channels_dict[res // 2] + use_fp16 = (res >= fp16_resolution) + block = DiscriminatorBlock(in_channels, tmp_channels, out_channels, resolution=res, + first_layer_idx=cur_layer_idx, use_fp16=use_fp16, **block_kwargs, **common_kwargs) + setattr(self, f'b{res}', block) + cur_layer_idx += block.num_layers + if c_dim > 0: + self.mapping = MappingNetwork(z_dim=0, c_dim=c_dim, w_dim=cmap_dim, num_ws=None, w_avg_beta=None, **mapping_kwargs) + self.b4 = DiscriminatorEpilogue(channels_dict[4], cmap_dim=cmap_dim, resolution=4, **epilogue_kwargs, **common_kwargs) + + def forward(self, img, c, update_emas=False, **block_kwargs): + _ = update_emas # unused + x = None + for res in self.block_resolutions: + block = getattr(self, f'b{res}') + x, img = block(x, img, **block_kwargs) + + cmap = None + if self.c_dim > 0: + cmap = self.mapping(None, c) + x = self.b4(x, img, cmap) + return x + + def extra_repr(self): + return f'c_dim={self.c_dim:d}, img_resolution={self.img_resolution:d}, img_channels={self.img_channels:d}' + +#---------------------------------------------------------------------------- \ No newline at end of file diff --git a/eg3d/training/superresolution.py b/eg3d/training/superresolution.py index 43321df2..5c4de2d6 100644 --- a/eg3d/training/superresolution.py +++ b/eg3d/training/superresolution.py @@ -21,7 +21,7 @@ import numpy as np from training.networks_stylegan3 import SynthesisLayer as AFSynthesisLayer - +from ipdb import set_trace as st #---------------------------------------------------------------------------- # for 512x512 generation diff --git a/eg3d/training/training_loop.py b/eg3d/training/training_loop.py index 63526bd9..ef2c21bf 100644 --- a/eg3d/training/training_loop.py +++ b/eg3d/training/training_loop.py @@ -15,6 +15,7 @@ import copy import json import pickle +import dill as pickle import psutil import PIL.Image import numpy as np @@ -29,7 +30,7 @@ from metrics import metric_main from camera_utils import LookAtPoseSampler from training.crosssection_utils import sample_cross_section - +from ipdb import set_trace as st #---------------------------------------------------------------------------- def setup_snapshot_image_grid(training_set, random_seed=0): @@ -67,8 +68,10 @@ def setup_snapshot_image_grid(training_set, random_seed=0): label_groups[label] = [indices[(i + gw) % len(indices)] for i in range(len(indices))] # Load data. - images, labels = zip(*[training_set[i] for i in grid_indices]) - return (gw, gh), np.stack(images), np.stack(labels) + images, labels, pc_arrays = zip(*[training_set[i] for i in grid_indices]) + # st() + # print('getitem types ---------------------->',type(images[0]), type(labels[0]), type(pc_arrays[0])) + return (gw, gh), np.stack(images), np.stack(labels), np.stack(pc_arrays) #---------------------------------------------------------------------------- @@ -156,9 +159,14 @@ def training_loop( print('Constructing networks...') common_kwargs = dict(c_dim=training_set.label_dim, img_resolution=training_set.resolution, img_channels=training_set.num_channels) G = dnnlib.util.construct_class_by_name(**G_kwargs, **common_kwargs).train().requires_grad_(False).to(device) # subclass of torch.nn.Module + # G = torch.nn.SyncBatchNorm.convert_sync_batchnorm(G).to(device) G.register_buffer('dataset_label_std', torch.tensor(training_set.get_label_std()).to(device)) + D = dnnlib.util.construct_class_by_name(**D_kwargs, **common_kwargs).train().requires_grad_(False).to(device) # subclass of torch.nn.Module + # D = torch.nn.SyncBatchNorm.convert_sync_batchnorm(D).to(device) G_ema = copy.deepcopy(G).eval() + # print('---------------------> G_kwargs', G_kwargs) + # print('---------------------> common_kwargs', common_kwargs) # Resume from existing pickle. if (resume_pkl is not None) and (rank == 0): @@ -166,13 +174,20 @@ def training_loop( with dnnlib.util.open_url(resume_pkl) as f: resume_data = legacy.load_network_pkl(f) for name, module in [('G', G), ('D', D), ('G_ema', G_ema)]: + # module = torch.nn.SyncBatchNorm.convert_sync_batchnorm(module).to(device) misc.copy_params_and_buffers(resume_data[name], module, require_all=False) # Print network summary tables. if rank == 0: z = torch.empty([batch_gpu, G.z_dim], device=device) c = torch.empty([batch_gpu, G.c_dim], device=device) - img = misc.print_module_summary(G, [z, c]) + + from training.volume import VolumeGenerator + if isinstance(G, VolumeGenerator): + pc = torch.empty([batch_gpu]+ [i for i in G.pc_dim], device=device) # (4, 1500, 9) + img = misc.print_module_summary(G, [z, c, pc]) + else: + img = misc.print_module_summary(G, [z, c]) misc.print_module_summary(D, [img, c]) # Setup augmentation. @@ -191,6 +206,7 @@ def training_loop( print(f'Distributing across {num_gpus} GPUs...') for module in [G, D, G_ema, augment_pipe]: if module is not None: + # module = torch.nn.SyncBatchNorm.convert_sync_batchnorm(module).to(device) for param in misc.params_and_buffers(module): if param.numel() > 0 and num_gpus > 1: torch.distributed.broadcast(param, src=0) @@ -202,13 +218,17 @@ def training_loop( phases = [] for name, module, opt_kwargs, reg_interval in [('G', G, G_opt_kwargs, G_reg_interval), ('D', D, D_opt_kwargs, D_reg_interval)]: if reg_interval is None: - opt = dnnlib.util.construct_class_by_name(params=module.parameters(), **opt_kwargs) # subclass of torch.optim.Optimizer + # print([name for name, p in module.named_parameters() if p.requires_grad]) # : empty + # opt = dnnlib.util.construct_class_by_name(params=[p for p in module.parameters() if p.requires_grad], **opt_kwargs) # subclass of torch.optim.Optimizer + opt = dnnlib.util.construct_class_by_name(params= module.parameters(), **opt_kwargs) # subclass of torch.optim.Optimizer phases += [dnnlib.EasyDict(name=name+'both', module=module, opt=opt, interval=1)] else: # Lazy regularization. mb_ratio = reg_interval / (reg_interval + 1) opt_kwargs = dnnlib.EasyDict(opt_kwargs) opt_kwargs.lr = opt_kwargs.lr * mb_ratio opt_kwargs.betas = [beta ** mb_ratio for beta in opt_kwargs.betas] + # print([name for name, p in module.named_parameters() if p.requires_grad]) # : empty + # opt = dnnlib.util.construct_class_by_name(params=[p for p in module.parameters() if p.requires_grad], **opt_kwargs) # subclass of torch.optim.Optimizer opt = dnnlib.util.construct_class_by_name(module.parameters(), **opt_kwargs) # subclass of torch.optim.Optimizer phases += [dnnlib.EasyDict(name=name+'main', module=module, opt=opt, interval=1)] phases += [dnnlib.EasyDict(name=name+'reg', module=module, opt=opt, interval=reg_interval)] @@ -218,17 +238,21 @@ def training_loop( if rank == 0: phase.start_event = torch.cuda.Event(enable_timing=True) phase.end_event = torch.cuda.Event(enable_timing=True) + print("Training phases:\n", [p.name for p in phases]) # Export sample images. grid_size = None grid_z = None grid_c = None + grid_pc = None if rank == 0: print('Exporting sample images...') - grid_size, images, labels = setup_snapshot_image_grid(training_set=training_set) + grid_size, images, labels, pointclouds = setup_snapshot_image_grid(training_set=training_set) + print("------------->", grid_size, images.shape, labels.shape, pointclouds.shape) save_image_grid(images, os.path.join(run_dir, 'reals.png'), drange=[0,255], grid_size=grid_size) grid_z = torch.randn([labels.shape[0], G.z_dim], device=device).split(batch_gpu) grid_c = torch.from_numpy(labels).to(device).split(batch_gpu) + grid_pc = torch.from_numpy(pointclouds).to(device).split(batch_gpu) # Initialize logs. if rank == 0: @@ -261,17 +285,23 @@ def training_loop( # Fetch training data. with torch.autograd.profiler.record_function('data_fetch'): - phase_real_img, phase_real_c = next(training_set_iterator) + phase_real_img, phase_real_c, phase_real_pc = next(training_set_iterator) phase_real_img = (phase_real_img.to(device).to(torch.float32) / 127.5 - 1).split(batch_gpu) phase_real_c = phase_real_c.to(device).split(batch_gpu) + phase_real_pc = phase_real_pc.to(device).split(batch_gpu) all_gen_z = torch.randn([len(phases) * batch_size, G.z_dim], device=device) all_gen_z = [phase_gen_z.split(batch_gpu) for phase_gen_z in all_gen_z.split(batch_size)] - all_gen_c = [training_set.get_label(np.random.randint(len(training_set))) for _ in range(len(phases) * batch_size)] + # same indices for c and pc + gen_indices = [np.random.randint(len(training_set)) for _ in range(len(phases) * batch_size)] + all_gen_c = [training_set.get_label(idx) for idx in gen_indices] all_gen_c = torch.from_numpy(np.stack(all_gen_c)).pin_memory().to(device) all_gen_c = [phase_gen_c.split(batch_gpu) for phase_gen_c in all_gen_c.split(batch_size)] + all_gen_pc = [training_set.get_pointcloud(idx) for idx in gen_indices] + all_gen_pc = torch.from_numpy(np.stack(all_gen_pc)).pin_memory().to(device) + all_gen_pc = [phase_gen_pc.split(batch_gpu) for phase_gen_pc in all_gen_pc.split(batch_size)] # Execute training phases. - for phase, phase_gen_z, phase_gen_c in zip(phases, all_gen_z, all_gen_c): + for phase, phase_gen_z, phase_gen_c, phase_gen_pc in zip(phases, all_gen_z, all_gen_c, all_gen_pc): if batch_idx % phase.interval != 0: continue if phase.start_event is not None: @@ -280,8 +310,12 @@ def training_loop( # Accumulate gradients. phase.opt.zero_grad(set_to_none=True) phase.module.requires_grad_(True) - for real_img, real_c, gen_z, gen_c in zip(phase_real_img, phase_real_c, phase_gen_z, phase_gen_c): - loss.accumulate_gradients(phase=phase.name, real_img=real_img, real_c=real_c, gen_z=gen_z, gen_c=gen_c, gain=phase.interval, cur_nimg=cur_nimg) + + # if rank == 0: + # print('############# current phase:', phase, '############') + + for real_img, real_c, gen_z, gen_c, gen_pc in zip(phase_real_img, phase_real_c, phase_gen_z, phase_gen_c, phase_gen_pc): + loss.accumulate_gradients(phase=phase.name, real_img=real_img, real_c=real_c, gen_z=gen_z, gen_c=gen_c, gen_pc=gen_pc, gain=phase.interval, cur_nimg=cur_nimg) phase.module.requires_grad_(False) # Update weights. @@ -328,6 +362,7 @@ def training_loop( # Perform maintenance tasks once per tick. done = (cur_nimg >= total_kimg * 1000) if (not done) and (cur_tick != 0) and (cur_nimg < tick_start_nimg + kimg_per_tick * 1000): + # if (not done) and (cur_tick != 0) and (cur_nimg < tick_start_nimg + kimg_per_tick * 2): continue # Print status line, accumulating the same information in training_stats. @@ -358,7 +393,12 @@ def training_loop( # Save image snapshot. if (rank == 0) and (image_snapshot_ticks is not None) and (done or cur_tick % image_snapshot_ticks == 0): - out = [G_ema(z=z, c=c, noise_mode='const') for z, c in zip(grid_z, grid_c)] + if all(x != 0 for x in pointclouds.shape): + # st() + out = [G_ema(z=z, c=c, pc=pc, noise_mode='const') for z, c, pc in zip(grid_z, grid_c, grid_pc)] + else: + st() + out = [G_ema(z=z, c=c, noise_mode='const') for z, c in zip(grid_z, grid_c)] images = torch.cat([o['image'].cpu() for o in out]).numpy() images_raw = torch.cat([o['image_raw'].cpu() for o in out]).numpy() images_depth = -torch.cat([o['image_depth'].cpu() for o in out]).numpy() @@ -399,7 +439,7 @@ def training_loop( for name, module in [('G', G), ('D', D), ('G_ema', G_ema), ('augment_pipe', augment_pipe)]: if module is not None: if num_gpus > 1: - misc.check_ddp_consistency(module, ignore_regex=r'.*\.[^.]+_(avg|ema)') + misc.check_ddp_consistency(module, ignore_regex=r'.*\.[^.]+_(avg|ema|mean)') module = copy.deepcopy(module).eval().requires_grad_(False).cpu() snapshot_data[name] = module del module # conserve memory @@ -410,6 +450,7 @@ def training_loop( # Evaluate metrics. if (snapshot_data is not None) and (len(metrics) > 0): + # if (snapshot_data is not None) and (len(metrics) > 0) and batch_idx % 10 ==0: if rank == 0: print(run_dir) print('Evaluating metrics...') diff --git a/eg3d/training/triplane.py b/eg3d/training/triplane.py index 1b7d48de..046e3565 100644 --- a/eg3d/training/triplane.py +++ b/eg3d/training/triplane.py @@ -102,6 +102,7 @@ def sample_mixed(self, coordinates, directions, ws, truncation_psi=1, truncation return self.renderer.run_model(planes, self.decoder, coordinates, directions, self.rendering_kwargs) def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, neural_rendering_resolution=None, update_emas=False, cache_backbone=False, use_cached_backbone=False, **synthesis_kwargs): + st() # Render a batch of generated images. ws = self.mapping(z, c, truncation_psi=truncation_psi, truncation_cutoff=truncation_cutoff, update_emas=update_emas) return self.synthesis(ws, c, update_emas=update_emas, neural_rendering_resolution=neural_rendering_resolution, cache_backbone=cache_backbone, use_cached_backbone=use_cached_backbone, **synthesis_kwargs) diff --git a/eg3d/training/volume.py b/eg3d/training/volume.py new file mode 100644 index 00000000..8817752a --- /dev/null +++ b/eg3d/training/volume.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +import torch +from torch_utils import persistence +# from training.networks_stylegan2 import Generator as StyleGAN2Backbone +# from training.networks_stylegan2 import FullyConnectedLayer + +# ### add 1d pc_ws to cur_ws, still tri-plane +# from training.networks_stylegan2_volume import Generator as VolumeBackbone +# from training.networks_stylegan2_volume import FullyConnectedLayer + +### no pc_ws, change snthesis_block to 3D, where output img is volume, and cat with pointcloud volume +from training.networks_stylegan2_syn_unet import Generator as VolumeBackbone +from training.networks_stylegan2_syn_unet import FullyConnectedLayer + +# from training.volumetric_rendering.renderer import ImportanceRenderer +from training.volumetric_rendering.renderer_volume import VolumeImportanceRenderer +from training.volumetric_rendering.ray_sampler import RaySampler +import dnnlib + +from ipdb import set_trace as st + +@persistence.persistent_class +class VolumeGenerator(torch.nn.Module): + def __init__(self, + z_dim, # Input latent (Z) dimensionality. + c_dim, # Conditioning label (C) dimensionality. + w_dim, # Intermediate latent (W) dimensionality. + img_resolution, # Output resolution. + img_channels, # Number of output color channels. + ####### newly added parameters ###### + pc_dim, # Conditioning poincloud (PC) dimensionality. + volume_res, # Volume resolution. + decoder_dim, + noise_strength, + ########################################## + sr_num_fp16_res = 0, + mapping_kwargs = {}, # Arguments for MappingNetwork. + rendering_kwargs = {}, + sr_kwargs = {}, + **synthesis_kwargs, # Arguments for SynthesisNetwork. + ): + super().__init__() + self.z_dim=z_dim + self.c_dim=c_dim + ####### newly added parameters ###### + self.pc_dim=pc_dim + self.volume_res=volume_res + ########################################## + self.w_dim=w_dim + self.img_resolution=img_resolution + self.img_channels=img_channels + # self.renderer = ImportanceRenderer() + self.renderer = VolumeImportanceRenderer() + self.ray_sampler = RaySampler() + ## ------ change backbone to 3d CONV Unet -------- + # self.backbone = StyleGAN2Backbone(z_dim, c_dim, w_dim, img_resolution=256, img_channels=32*3, mapping_kwargs=mapping_kwargs, **synthesis_kwargs) + self.backbone = VolumeBackbone(z_dim, c_dim, w_dim, \ + pc_dim=pc_dim, volume_res=volume_res, noise_strength=noise_strength,\ + img_resolution=256, img_channels=32*3, mapping_kwargs=mapping_kwargs, **synthesis_kwargs) + ## + self.superresolution = dnnlib.util.construct_class_by_name(class_name=rendering_kwargs['superresolution_module'], channels=32, img_resolution=img_resolution, sr_num_fp16_res=sr_num_fp16_res, sr_antialias=rendering_kwargs['sr_antialias'], **sr_kwargs) + # self.decoder = OSGDecoder(32, {'decoder_lr_mul': rendering_kwargs.get('decoder_lr_mul', 1), 'decoder_output_dim': 32}) + self.decoder = OSGDecoder(decoder_dim, \ + {'decoder_lr_mul': rendering_kwargs.get('decoder_lr_mul', 1), \ + 'decoder_output_dim': 32, \ + 'use_ray_directions': rendering_kwargs.get('use_ray_directions', False)}) # input_dim=8 for volume + self.neural_rendering_resolution = 64 + self.rendering_kwargs = rendering_kwargs + + self._last_planes = None + self.log_idx = 0 + + def mapping(self, z, c, truncation_psi=1, truncation_cutoff=None, update_emas=False): + if self.rendering_kwargs['c_gen_conditioning_zero']: # True + c = torch.zeros_like(c) + # st() + else: + st() # make the generation condition on camera pose + return self.backbone.mapping(z, c * self.rendering_kwargs.get('c_scale', 0), truncation_psi=truncation_psi, truncation_cutoff=truncation_cutoff, update_emas=update_emas) + + def synthesis(self, ws, c, pc=None, neural_rendering_resolution=None, update_emas=False, cache_backbone=False, use_cached_backbone=False, **synthesis_kwargs): + cam2world_matrix = c[:, :16].view(-1, 4, 4) + intrinsics = c[:, 16:25].view(-1, 3, 3) + + if neural_rendering_resolution is None: + neural_rendering_resolution = self.neural_rendering_resolution + else: + self.neural_rendering_resolution = neural_rendering_resolution + + # Create a batch of rays for volume rendering + ray_origins, ray_directions = self.ray_sampler(cam2world_matrix, intrinsics, neural_rendering_resolution) + # TODO Oct 16: check aligning result with lego tensorf in mvsnerf + + # Create triplanes by running StyleGAN backbone + N, M, _ = ray_origins.shape + if use_cached_backbone and self._last_planes is not None: + planes = self._last_planes + st() # assert not coming into this block + else: + planes = self.backbone.synthesis(ws, pc=pc, box_warp=self.rendering_kwargs['box_warp'], update_emas=update_emas, **synthesis_kwargs) + # this will call: SynthesisNetwork.forward() + + if cache_backbone: + st() # assert not coming into this block + self._last_planes = planes + + # Reshape output into three 32-channel planes + if isinstance(planes, tuple): + planes = list(planes) + planes[0]=planes[0].view(len(planes[0]), 3, 32, planes[0].shape[-2], planes[0].shape[-1]) + # st() + else: + try: + planes = planes.view(len(planes), 3, 32, planes.shape[-2], planes.shape[-1]) + except: + # TODO: replace with volume: + # 1. no reshape + # 2. . + # (do nothing) + # st() + pass + + # Perform volume rendering + ## already adapted to volume + # st() + feature_samples, depth_samples, weights_samples = self.renderer(planes, self.decoder, ray_origins, ray_directions, self.rendering_kwargs) # channels last + # st() + + + # Reshape into 'raw' neural-rendered image + H = W = self.neural_rendering_resolution + feature_image = feature_samples.permute(0, 2, 1).reshape(N, feature_samples.shape[-1], H, W).contiguous() + depth_image = depth_samples.permute(0, 2, 1).reshape(N, 1, H, W) + + # Run superresolution to get final image + rgb_image = feature_image[:, :3] + # st() + sr_image = self.superresolution(rgb_image, feature_image, ws, noise_mode=self.rendering_kwargs['superresolution_noise_mode'], **{k:synthesis_kwargs[k] for k in synthesis_kwargs.keys() if k != 'noise_mode'}) + + return {'image': sr_image, 'image_raw': rgb_image, 'image_depth': depth_image} + + def sample(self, coordinates, directions, z, c, truncation_psi=1, truncation_cutoff=None, update_emas=False, **synthesis_kwargs): + # Compute RGB features, density for arbitrary 3D coordinates. Mostly used for extracting shapes. + ws = self.mapping(z, c, truncation_psi=truncation_psi, truncation_cutoff=truncation_cutoff, update_emas=update_emas) + planes = self.backbone.synthesis(ws, pc=pc, box_warp=box_warp, update_emas=update_emas, **synthesis_kwargs) + planes = planes.view(len(planes), 3, 32, planes.shape[-2], planes.shape[-1]) + return self.renderer.run_model(planes, self.decoder, coordinates, directions, self.rendering_kwargs) + + def sample_mixed(self, coordinates, directions, ws, pc=None, box_warp=None, truncation_psi=1, truncation_cutoff=None, update_emas=False, **synthesis_kwargs): + # Same as sample, but expects latent vectors 'ws' instead of Gaussian noise 'z' + planes = self.backbone.synthesis(ws, pc=pc, box_warp=box_warp, update_emas = update_emas, **synthesis_kwargs) + if isinstance(planes, tuple): + planes = list(planes) + planes[0]=planes[0].view(len(planes[0]), 3, 32, planes[0].shape[-2], planes[0].shape[-1]) + # st() + elif planes.shape[-1]!=planes.shape[-3]: + planes = planes.view(len(planes), 3, 32, planes.shape[-2], planes.shape[-1]) + return self.renderer.run_model(planes, self.decoder, coordinates, directions, self.rendering_kwargs) + + def forward(self, z, c, pc, truncation_psi=1, truncation_cutoff=None, neural_rendering_resolution=None, update_emas=False, cache_backbone=False, use_cached_backbone=False, **synthesis_kwargs): + if pc.shape[-2:] != (1024,9): + st() + + # self.log_idx= self.log_idx +1 + # print('foward in Volumegenerator', self.log_idx) + + # Render a batch of generated images. + ws = self.mapping(z, c, truncation_psi=truncation_psi, truncation_cutoff=truncation_cutoff, update_emas=update_emas) # (4, 14, 512) + return self.synthesis(ws, c, pc=pc, update_emas=update_emas, neural_rendering_resolution=neural_rendering_resolution, cache_backbone=cache_backbone, use_cached_backbone=use_cached_backbone, **synthesis_kwargs) + + + + +class OSGDecoder(torch.nn.Module): + def __init__(self, n_features, options): + super().__init__() + self.hidden_dim = 64 + # if n_features != 8: + # st() + + self.use_ray_directions = options['use_ray_directions'] + if self.use_ray_directions: + n_features += 3 + + + self.net = torch.nn.Sequential( + FullyConnectedLayer(n_features, self.hidden_dim, lr_multiplier=options['decoder_lr_mul']), + torch.nn.Softplus(), + FullyConnectedLayer(self.hidden_dim, 1 + options['decoder_output_dim'], lr_multiplier=options['decoder_lr_mul']) + ) + + + def forward(self, sampled_features, ray_directions): + # st() # x.shape + # Aggregate features + + sampled_features = sampled_features.mean(1) # tri-plane: mean of 3 planes; volume: only one volume, so mean() is the same as squeeze + + if self.use_ray_directions: + sampled_features = torch.cat([sampled_features, ray_directions], -1) + + x = sampled_features + + N, M, C = x.shape + x = x.view(N*M, C) + + x = self.net(x) + x = x.view(N, M, -1) + + rgb = torch.sigmoid(x[..., 1:])*(1 + 2*0.001) - 0.001 # Uses sigmoid clamping from MipNeRF + sigma = x[..., 0:1] + + return {'rgb': rgb, 'sigma': sigma} diff --git a/eg3d/training/volumetric_rendering/ray_marcher.py b/eg3d/training/volumetric_rendering/ray_marcher.py index c2c427f7..c10b7591 100644 --- a/eg3d/training/volumetric_rendering/ray_marcher.py +++ b/eg3d/training/volumetric_rendering/ray_marcher.py @@ -16,6 +16,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from ipdb import set_trace as st class MipRayMarcher2(nn.Module): def __init__(self): @@ -24,6 +25,7 @@ def __init__(self): def run_forward(self, colors, densities, depths, rendering_options): deltas = depths[:, :, 1:] - depths[:, :, :-1] + # st() colors_mid = (colors[:, :, :-1] + colors[:, :, 1:]) / 2 densities_mid = (densities[:, :, :-1] + densities[:, :, 1:]) / 2 depths_mid = (depths[:, :, :-1] + depths[:, :, 1:]) / 2 @@ -49,7 +51,9 @@ def run_forward(self, colors, densities, depths, rendering_options): composite_depth = torch.nan_to_num(composite_depth, float('inf')) composite_depth = torch.clamp(composite_depth, torch.min(depths), torch.max(depths)) + # st() if rendering_options.get('white_back', False): + # if rendering_options.get('white_back') == True: composite_rgb = composite_rgb + 1 - weight_total composite_rgb = composite_rgb * 2 - 1 # Scale to (-1, 1) diff --git a/eg3d/training/volumetric_rendering/ray_sampler.py b/eg3d/training/volumetric_rendering/ray_sampler.py index 00dd07b9..3cd59726 100644 --- a/eg3d/training/volumetric_rendering/ray_sampler.py +++ b/eg3d/training/volumetric_rendering/ray_sampler.py @@ -14,6 +14,7 @@ """ import torch +from ipdb import set_trace as st class RaySampler(torch.nn.Module): def __init__(self): @@ -40,6 +41,8 @@ def forward(self, cam2world_matrix, intrinsics, resolution): cy = intrinsics[:, 1, 2] sk = intrinsics[:, 0, 1] + # st() + uv = torch.stack(torch.meshgrid(torch.arange(resolution, dtype=torch.float32, device=cam2world_matrix.device), torch.arange(resolution, dtype=torch.float32, device=cam2world_matrix.device), indexing='ij')) * (1./resolution) + (0.5/resolution) uv = uv.flip(0).reshape(2, -1).transpose(1, 0) uv = uv.unsqueeze(0).repeat(cam2world_matrix.shape[0], 1, 1) @@ -47,6 +50,8 @@ def forward(self, cam2world_matrix, intrinsics, resolution): x_cam = uv[:, :, 0].view(N, -1) y_cam = uv[:, :, 1].view(N, -1) z_cam = torch.ones((N, M), device=cam2world_matrix.device) + # if not torch.all(intrinsics==0): + # st() x_lift = (x_cam - cx.unsqueeze(-1) + cy.unsqueeze(-1)*sk.unsqueeze(-1)/fy.unsqueeze(-1) - sk.unsqueeze(-1)*y_cam/fy.unsqueeze(-1)) / fx.unsqueeze(-1) * z_cam y_lift = (y_cam - cy.unsqueeze(-1)) / fy.unsqueeze(-1) * z_cam diff --git a/eg3d/training/volumetric_rendering/renderer_volume.py b/eg3d/training/volumetric_rendering/renderer_volume.py new file mode 100644 index 00000000..a22db10a --- /dev/null +++ b/eg3d/training/volumetric_rendering/renderer_volume.py @@ -0,0 +1,307 @@ +# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +""" +The renderer is a module that takes in rays, decides where to sample along each +ray, and computes pixel colors using the volume rendering equation. +""" + +import math +import torch +import torch.nn as nn +import torch.nn.functional as F + +from training.volumetric_rendering.ray_marcher import MipRayMarcher2 +from training.volumetric_rendering import math_utils + +from ipdb import set_trace as st + +def generate_planes(): + """ + Defines planes by the three vectors that form the "axes" of the + plane. Should work with arbitrary number of planes and planes of + arbitrary orientation. + """ + return torch.tensor([[[1, 0, 0], + [0, 1, 0], + [0, 0, 1]], + [[1, 0, 0], + [0, 0, 1], + [0, 1, 0]], + [[0, 0, 1], + [1, 0, 0], + [0, 1, 0]]], dtype=torch.float32) + +def project_onto_planes(planes, coordinates): + """ + Does a projection of a 3D point onto a batch of 2D planes, + returning 2D plane coordinates. + + Takes plane axes of shape n_planes, 3, 3 + # Takes coordinates of shape N, M, 3 + # returns projections of shape N*n_planes, M, 2 + """ + N, M, C = coordinates.shape + n_planes, _, _ = planes.shape + coordinates = coordinates.unsqueeze(1).expand(-1, n_planes, -1, -1).reshape(N*n_planes, M, 3) + inv_planes = torch.linalg.inv(planes).unsqueeze(0).expand(N, -1, -1, -1).reshape(N*n_planes, 3, 3) + projections = torch.bmm(coordinates, inv_planes) + return projections[..., :2] + +def sample_from_planes(plane_axes, plane_features, coordinates, mode='bilinear', padding_mode='zeros', box_warp=None): + assert padding_mode == 'zeros' + N, n_planes, C, H, W = plane_features.shape + _, M, _ = coordinates.shape + plane_features = plane_features.view(N*n_planes, C, H, W) + + ## box_warp: the covered area of tri-plane/volume + ## so the sampled coordinated will be normalized to the tri-plane/volume here, no need to process it like in the get_rays_ndc in the mvsnerf + coordinates = (2/box_warp) * coordinates # TODO: add specific box bounds + + projected_coordinates = project_onto_planes(plane_axes, coordinates).unsqueeze(1) + + output_features = torch.nn.functional.grid_sample(plane_features, projected_coordinates.float(), mode=mode, padding_mode=padding_mode, align_corners=False).permute(0, 3, 2, 1).reshape(N, n_planes, M, C) + return output_features + +def sample_from_volume(volume_features, coordinates, mode='bilinear', padding_mode='zeros', box_warp=None): + + #### MVSNerf gird sample as below: + # input_feat = index_point_feature_batch(volume_feature, rays_ndc, grid_sample_mode=grid_sample_mode) if torch.is_tensor(volume_feature) else volume_feature(rays_ndc) + # features = F.grid_sample(volume_feature, grid, align_corners=False, mode=grid_sample_mode)[:,:,0].permute(0,2,3,1) # mode='nearest' + + assert padding_mode == 'zeros' + N, C, D, H, W = volume_features.shape + # st() + _, M, _ = coordinates.shape # N,M,3 where N is batch size + n_planes = 1 + # plane_features = plane_features.view(N*n_planes, C, H, W) + + ## box_warp: the covered area of tri-plane/volume + ## so the sampled coordinated will be normalized to the tri-plane/volume here, no need to process it like in the get_rays_ndc in the mvsnerf + coordinates = (2/box_warp) * coordinates # -1~1 + coordinates = coordinates[:, None, None, ...] # shape: (B,1,1,M,3) to correctly use 3D grid sampled + + # projected_coordinates = project_onto_planes(plane_axes, coordinates).unsqueeze(1) + output_features = torch.nn.functional.grid_sample(volume_features, coordinates.float(), mode=mode, padding_mode=padding_mode, align_corners=False)[:,:,0] #.permute(0, 2, 3, 1).reshape(N, n_planes, M, C) + output_features = output_features.permute(0, 2, 3, 1).reshape(N, n_planes, M, C) + return output_features + +def sample_from_3dgrid(grid, coordinates): + """ + Expects coordinates in shape (batch_size, num_points_per_batch, 3) + Expects grid in shape (1, channels, H, W, D) + (Also works if grid has batch size) + Returns sampled features of shape (batch_size, num_points_per_batch, feature_channels) + """ + batch_size, n_coords, n_dims = coordinates.shape + sampled_features = torch.nn.functional.grid_sample(grid.expand(batch_size, -1, -1, -1, -1), + coordinates.reshape(batch_size, 1, 1, -1, n_dims), + mode='bilinear', padding_mode='zeros', align_corners=False) + N, C, H, W, D = sampled_features.shape + sampled_features = sampled_features.permute(0, 4, 3, 2, 1).reshape(N, H*W*D, C) + return sampled_features + +class VolumeImportanceRenderer(torch.nn.Module): + def __init__(self): + super().__init__() + self.ray_marcher = MipRayMarcher2() + self.plane_axes = generate_planes() + + def forward(self, planes, decoder, ray_origins, ray_directions, rendering_options): + # planes.shape: torch.Size([4, 8, 64, 64, 64]) + self.plane_axes = self.plane_axes.to(ray_origins.device) + + if rendering_options['ray_start'] == rendering_options['ray_end'] == 'auto': + ray_start, ray_end = math_utils.get_ray_limits_box(ray_origins, ray_directions, box_side_length=rendering_options['box_warp']) + is_ray_valid = ray_end > ray_start + if torch.any(is_ray_valid).item(): + ray_start[~is_ray_valid] = ray_start[is_ray_valid].min() + ray_end[~is_ray_valid] = ray_start[is_ray_valid].max() + depths_coarse = self.sample_stratified(ray_origins, ray_start, ray_end, rendering_options['depth_resolution'], rendering_options['disparity_space_sampling']) + else: + # Create stratified depth samples + depths_coarse = self.sample_stratified(ray_origins, rendering_options['ray_start'], rendering_options['ray_end'], rendering_options['depth_resolution'], rendering_options['disparity_space_sampling']) + + batch_size, num_rays, samples_per_ray, _ = depths_coarse.shape + + # Coarse Pass + sample_coordinates = (ray_origins.unsqueeze(-2) + depths_coarse * ray_directions.unsqueeze(-2)).reshape(batch_size, -1, 3) + sample_directions = ray_directions.unsqueeze(-2).expand(-1, -1, samples_per_ray, -1).reshape(batch_size, -1, 3) + + + out = self.run_model(planes, decoder, sample_coordinates, sample_directions, rendering_options) + colors_coarse = out['rgb'] + + densities_coarse = out['sigma'] + colors_coarse = colors_coarse.reshape(batch_size, num_rays, samples_per_ray, colors_coarse.shape[-1]) + densities_coarse = densities_coarse.reshape(batch_size, num_rays, samples_per_ray, 1) + + # Fine Pass + # st() + N_importance = rendering_options['depth_resolution_importance'] + if N_importance > 0: + _, _, weights = self.ray_marcher(colors_coarse, densities_coarse, depths_coarse, rendering_options) + + depths_fine = self.sample_importance(depths_coarse, weights, N_importance) + + sample_directions = ray_directions.unsqueeze(-2).expand(-1, -1, N_importance, -1).reshape(batch_size, -1, 3) + sample_coordinates = (ray_origins.unsqueeze(-2) + depths_fine * ray_directions.unsqueeze(-2)).reshape(batch_size, -1, 3) + + out = self.run_model(planes, decoder, sample_coordinates, sample_directions, rendering_options) + colors_fine = out['rgb'] + densities_fine = out['sigma'] + colors_fine = colors_fine.reshape(batch_size, num_rays, N_importance, colors_fine.shape[-1]) + densities_fine = densities_fine.reshape(batch_size, num_rays, N_importance, 1) + + all_depths, all_colors, all_densities = self.unify_samples(depths_coarse, colors_coarse, densities_coarse, + depths_fine, colors_fine, densities_fine) + + # Aggregate + rgb_final, depth_final, weights = self.ray_marcher(all_colors, all_densities, all_depths, rendering_options) + else: + rgb_final, depth_final, weights = self.ray_marcher(colors_coarse, densities_coarse, depths_coarse, rendering_options) + + + return rgb_final, depth_final, weights.sum(2) + + def run_model(self, planes, decoder, sample_coordinates, sample_directions, options): + # st() + # if planes.shape[1]==96: + if isinstance(planes, list): + # st() + sampled_features_plane = sample_from_planes(self.plane_axes, planes[0], sample_coordinates, padding_mode='zeros', box_warp=options['box_warp']) + sampled_features_volume = sample_from_volume(planes[1], sample_coordinates, padding_mode='zeros', box_warp=options['box_warp']) + + mask = torch.ones_like(sampled_features_plane) + mask[...,:sampled_features_volume.shape[-1]]*=0 + p1d=(0,sampled_features_plane.shape[-1]-sampled_features_volume.shape[-1]) + sampled_features_volume = F.pad(sampled_features_volume, p1d, 'constant', 0) + sampled_features = mask*sampled_features_plane + (1-mask)*sampled_features_volume + # st() + + elif planes.shape[-3]!=planes.shape[-1]: + st() + sampled_features = sample_from_planes(self.plane_axes, planes, sample_coordinates, padding_mode='zeros', box_warp=options['box_warp']) + else: + # st() + sampled_features = sample_from_volume(planes, sample_coordinates, padding_mode='zeros', box_warp=options['box_warp']) + # st() # align sampled_features.shape and sampled_features.shape + + out = decoder(sampled_features, sample_directions) + # st() + if options.get('density_noise', 0) > 0: + out['sigma'] += torch.randn_like(out['sigma']) * options['density_noise'] + return out + + def sort_samples(self, all_depths, all_colors, all_densities): + _, indices = torch.sort(all_depths, dim=-2) + all_depths = torch.gather(all_depths, -2, indices) + all_colors = torch.gather(all_colors, -2, indices.expand(-1, -1, -1, all_colors.shape[-1])) + all_densities = torch.gather(all_densities, -2, indices.expand(-1, -1, -1, 1)) + return all_depths, all_colors, all_densities + + def unify_samples(self, depths1, colors1, densities1, depths2, colors2, densities2): + all_depths = torch.cat([depths1, depths2], dim = -2) + all_colors = torch.cat([colors1, colors2], dim = -2) + all_densities = torch.cat([densities1, densities2], dim = -2) + + _, indices = torch.sort(all_depths, dim=-2) + all_depths = torch.gather(all_depths, -2, indices) + all_colors = torch.gather(all_colors, -2, indices.expand(-1, -1, -1, all_colors.shape[-1])) + all_densities = torch.gather(all_densities, -2, indices.expand(-1, -1, -1, 1)) + + return all_depths, all_colors, all_densities + + def sample_stratified(self, ray_origins, ray_start, ray_end, depth_resolution, disparity_space_sampling=False): + """ + Return depths of approximately uniformly spaced samples along rays. + """ + N, M, _ = ray_origins.shape + if disparity_space_sampling: + depths_coarse = torch.linspace(0, + 1, + depth_resolution, + device=ray_origins.device).reshape(1, 1, depth_resolution, 1).repeat(N, M, 1, 1) + depth_delta = 1/(depth_resolution - 1) + depths_coarse += torch.rand_like(depths_coarse) * depth_delta + depths_coarse = 1./(1./ray_start * (1. - depths_coarse) + 1./ray_end * depths_coarse) + else: + if type(ray_start) == torch.Tensor: + depths_coarse = math_utils.linspace(ray_start, ray_end, depth_resolution).permute(1,2,0,3) + depth_delta = (ray_end - ray_start) / (depth_resolution - 1) + depths_coarse += torch.rand_like(depths_coarse) * depth_delta[..., None] + else: + depths_coarse = torch.linspace(ray_start, ray_end, depth_resolution, device=ray_origins.device).reshape(1, 1, depth_resolution, 1).repeat(N, M, 1, 1) + depth_delta = (ray_end - ray_start)/(depth_resolution - 1) + depths_coarse += torch.rand_like(depths_coarse) * depth_delta + + return depths_coarse + + def sample_importance(self, z_vals, weights, N_importance): + """ + Return depths of importance sampled points along rays. See NeRF importance sampling for more. + """ + with torch.no_grad(): + batch_size, num_rays, samples_per_ray, _ = z_vals.shape + + z_vals = z_vals.reshape(batch_size * num_rays, samples_per_ray) + weights = weights.reshape(batch_size * num_rays, -1) # -1 to account for loss of 1 sample in MipRayMarcher + + # smooth weights + weights = torch.nn.functional.max_pool1d(weights.unsqueeze(1).float(), 2, 1, padding=1) + weights = torch.nn.functional.avg_pool1d(weights, 2, 1).squeeze() + weights = weights + 0.01 + + z_vals_mid = 0.5 * (z_vals[: ,:-1] + z_vals[: ,1:]) + importance_z_vals = self.sample_pdf(z_vals_mid, weights[:, 1:-1], + N_importance).detach().reshape(batch_size, num_rays, N_importance, 1) + return importance_z_vals + + def sample_pdf(self, bins, weights, N_importance, det=False, eps=1e-5): + """ + Sample @N_importance samples from @bins with distribution defined by @weights. + Inputs: + bins: (N_rays, N_samples_+1) where N_samples_ is "the number of coarse samples per ray - 2" + weights: (N_rays, N_samples_) + N_importance: the number of samples to draw from the distribution + det: deterministic or not + eps: a small number to prevent division by zero + Outputs: + samples: the sampled samples + """ + N_rays, N_samples_ = weights.shape + weights = weights + eps # prevent division by zero (don't do inplace op!) + pdf = weights / torch.sum(weights, -1, keepdim=True) # (N_rays, N_samples_) + cdf = torch.cumsum(pdf, -1) # (N_rays, N_samples), cumulative distribution function + cdf = torch.cat([torch.zeros_like(cdf[: ,:1]), cdf], -1) # (N_rays, N_samples_+1) + # padded to 0~1 inclusive + + if det: + u = torch.linspace(0, 1, N_importance, device=bins.device) + u = u.expand(N_rays, N_importance) + else: + u = torch.rand(N_rays, N_importance, device=bins.device) + u = u.contiguous() + + inds = torch.searchsorted(cdf, u, right=True) + below = torch.clamp_min(inds-1, 0) + above = torch.clamp_max(inds, N_samples_) + + inds_sampled = torch.stack([below, above], -1).view(N_rays, 2*N_importance) + cdf_g = torch.gather(cdf, 1, inds_sampled).view(N_rays, N_importance, 2) + # st() + bins_g = torch.gather(bins, 1, inds_sampled).view(N_rays, N_importance, 2) + + denom = cdf_g[...,1]-cdf_g[...,0] + denom[denom