├── .dockerignore ├── .gitignore ├── Dockerfile ├── GGXBRDF ├── ggx_brdf.py └── onb.py ├── LICENSE ├── assets ├── sorted_teaser_low_res.png └── sorted_teaser_new.jpg ├── data_processing ├── camera_related │ └── camera_config.py ├── extract_measurements │ ├── extract.py │ └── run.sh ├── finetune │ ├── AUTO_lumitexel_net_inference.py │ ├── AUTO_planar_scanner_net_inference.py │ ├── choose_branch_net.py │ ├── finetune_pass1.py │ ├── finetune_pass2.py │ ├── finetune_pass3.py │ ├── fitting_master.py │ ├── gather_finetune_results.py │ ├── infer.py │ ├── latent_net.py │ ├── planar_scanner_net_inference.py │ ├── run.sh │ └── split_data_and_prepare_for_server.py ├── fitting │ ├── fit_latent.py │ ├── latent_controller.py │ ├── latent_dataset.py │ ├── latent_mlp.py │ ├── latent_solver.py │ ├── mlp_model.py │ ├── run.sh │ ├── texture.py │ └── utils.py ├── generate_uv │ ├── extract_plane.py │ ├── run.sh │ └── warp.py ├── parallel-patchmatch-master │ ├── main.py │ ├── patchmatch.py │ ├── run.sh │ ├── sift_flow_torch.py │ └── third_party │ │ └── flowiz │ │ ├── LICENSE │ │ ├── README.md │ │ └── flowiz.py ├── torch_renderer │ ├── camera.py │ ├── materials │ │ ├── __init__.py │ │ ├── ggx_brdf.py │ │ └── utils.py │ ├── onb.py │ ├── ray_trace.py │ ├── render_utils.py │ ├── setup_config.py │ ├── torch_render.py │ └── wallet_of_torch_renderer │ │ └── lightstage │ │ ├── cam_pos.bin │ │ ├── color_tensor.bin │ │ ├── extrinsic.yml │ │ ├── intrinsic.yml │ │ ├── lights.bin │ │ ├── mask.npy │ │ └── visualize_config_torch.bin └── visualize_results │ ├── rendering.py │ └── run.sh ├── database_data └── download.md ├── database_model └── download.md ├── download.sh ├── environment.yml ├── readme.md ├── requirements.txt └── run.sh /.dockerignore: -------------------------------------------------------------------------------- 1 | database_data/ -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | .hypothesis/ 51 | .pytest_cache/ 52 | 53 | # Translations 54 | *.mo 55 | *.pot 56 | 57 | # Django stuff: 58 | *.log 59 | local_settings.py 60 | db.sqlite3 61 | db.sqlite3-journal 62 | 63 | # Flask stuff: 64 | instance/ 65 | .webassets-cache 66 | 67 | # Scrapy stuff: 68 | .scrapy 69 | 70 | # Sphinx documentation 71 | docs/_build/ 72 | 73 | # PyBuilder 74 | target/ 75 | 76 | # Jupyter Notebook 77 | .ipynb_checkpoints 78 | 79 | # IPython 80 | profile_default/ 81 | ipython_config.py 82 | 83 | # pyenv 84 | .python-version 85 | 86 | # pipenv 87 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 88 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 89 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 90 | # install all needed dependencies. 91 | #Pipfile.lock 92 | 93 | # celery beat schedule file 94 | celerybeat-schedule 95 | 96 | # SageMath parsed files 97 | *.sage.py 98 | 99 | # Environments 100 | .env 101 | .venv 102 | env/ 103 | venv/ 104 | ENV/ 105 | env.bak/ 106 | venv.bak/ 107 | 108 | # Spyder project settings 109 | .spyderproject 110 | .spyproject 111 | 112 | # Rope project settings 113 | .ropeproject 114 | 115 | # mkdocs documentation 116 | /site 117 | 118 | # mypy 119 | .mypy_cache/ 120 | .dmypy.json 121 | dmypy.json 122 | 123 | # Pyre type checker 124 | .pyre/ 125 | device_configuration/ 126 | 127 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | FROM nvidia/cuda:12.1.0-devel-ubuntu22.04 2 | 3 | RUN apt-get update && \ 4 | apt-get install -y python3 curl libgl1 libglib2.0-0 5 | 6 | RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python3 7 | 8 | COPY . /app 9 | WORKDIR /app 10 | 11 | RUN pip install -r /app/requirements.txt 12 | 13 | RUN ln -s /usr/bin/python3 /usr/bin/python 14 | -------------------------------------------------------------------------------- /GGXBRDF/ggx_brdf.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import torch.nn.functional as F 4 | from onb import ONB 5 | 6 | 7 | class GGX_BRDF(object): 8 | 9 | @classmethod 10 | def eval( 11 | cls, 12 | local_onb: ONB, 13 | wi: torch.Tensor, 14 | wo: torch.Tensor, 15 | ax: torch.Tensor, 16 | ay: torch.Tensor, 17 | pd: torch.Tensor, 18 | ps: torch.Tensor, 19 | is_diff: torch.Tensor = None, 20 | specular_component: str = "D_F_G_B", 21 | ) -> torch.Tensor: 22 | """ 23 | Evaluate brdf in shading coordinate(ntb space). 24 | 25 | For each shading point, we will build a coordinate system to calculate 26 | the brdf. The coordinate is usually expressed as ntb. n is the normal 27 | of the shading point. 28 | 29 | Args: 30 | local_onb: a coordinate system to do shade 31 | wi: incident light in world space, of shape (batch, lightnum, 3) 32 | wo: outgoing light in world space, of shape (batch, 3) 33 | ax: shape (batch, 1) 34 | ay: shape (batch, 1) 35 | pd: shape (batch, channel), range [0, 1] 36 | ps: shape (batch, channel), range [0, 10] 37 | is_diff: shape (batch, lightnum), if the value is 0, eval "pd_only", else if 38 | the value is 1, eval "ps_only". If `is_diff` is None, means "both" 39 | specular_component: the ingredient of BRDF, usually "D_F_G_B", B means bottom 40 | 41 | Returns: 42 | brdf: (batch, lightnum, channel) 43 | meta: 44 | """ 45 | N = wi.size(1) 46 | 47 | # transform wi and wo to local frame 48 | wi_local = local_onb.transform(wi) # (batch, lightnum, 3) 49 | wo_local = local_onb.transform(wo) # (batch, 3) 50 | 51 | meta = {} 52 | 53 | a = torch.unsqueeze(pd / math.pi, dim=1) # (batch, 1, 1) 54 | b = cls.ggx_brdf_aniso(wi_local, wo_local, ax, ay, specular_component) # (batch, lightnum, 1) 55 | ps = torch.unsqueeze(ps, dim=1) # (batch, 1, channel) 56 | 57 | if is_diff is None: 58 | brdf = a + b * ps 59 | else: 60 | is_diff_ = is_diff.unsqueeze(2) 61 | brdf = a.repeat(1, N, 1) * (1 - is_diff_) + b * ps * is_diff_ 62 | 63 | meta['pd'] = a 64 | meta['ps'] = b * ps 65 | 66 | return brdf, meta 67 | 68 | @classmethod 69 | def ggx_brdf_aniso( 70 | cls, 71 | wi: torch.Tensor, 72 | wo: torch.Tensor, 73 | ax: torch.Tensor, 74 | ay: torch.Tensor, 75 | specular_component: str 76 | ) -> torch.Tensor: 77 | """ 78 | Calculate anisotropy ggx brdf in shading coordinate. 79 | 80 | Args: 81 | wi: incident light in ntb space, of shape (batch, lightnum, 3) 82 | wo: emergent light in ntb space, of shape (batch, 3) 83 | ax: shape (batch, 1) 84 | ay: shape (batch, 1) 85 | specular_component: the ingredient of BRDF, usually "D_F_G_B" 86 | 87 | Returns: 88 | brdf: shape (batch, lightnum, 1) 89 | """ 90 | 91 | lightnum = wi.size(1) 92 | wo = torch.unsqueeze(wo, dim=1).repeat(1, lightnum, 1) 93 | 94 | wi_z = wi[:, :, [2]] # (batch, lightnum, 1) 95 | wo_z = wo[:, :, [2]] 96 | denom = 4 * wi_z * wo_z # (batch, lightnum, 1) 97 | vhalf = F.normalize(wi + wo, dim=2) # (batch, lightnum, 3) 98 | 99 | # F 100 | tmp = torch.clamp(1.0 - torch.sum(wi * vhalf, dim=2, keepdim=True), 0, 1) 101 | F0 = 0.04 102 | Fresnel = F0 + (1 - F0) * tmp * tmp * tmp * tmp * tmp 103 | 104 | # D 105 | axayaz = torch.unsqueeze(torch.cat([ax, ay, torch.ones_like(ax)], dim=1), 106 | dim=1) # (batch, 1, 3) 107 | vhalf = vhalf / (axayaz + 1e-6) # (batch, lightnum, 3) 108 | vhalf_norm = torch.norm(vhalf, dim=2, keepdim=True) 109 | length = vhalf_norm * vhalf_norm # (batch, lightnum, 1) 110 | D = 1.0 / (math.pi * torch.unsqueeze(ax, dim=1) * 111 | torch.unsqueeze(ay, dim=1) * length * length) 112 | 113 | # G 114 | G = cls.ggx_G1_aniso(wi, ax, ay, wi_z) * cls.ggx_G1_aniso(wo, ax, ay, wo_z) 115 | 116 | tmp = torch.ones_like(denom) 117 | if "D" in specular_component: 118 | tmp = tmp * D 119 | if "F" in specular_component: 120 | tmp = tmp * Fresnel 121 | if "G" in specular_component: 122 | tmp = tmp * G 123 | if "B" in specular_component: 124 | tmp = tmp / (denom + 1e-6) 125 | 126 | # some samples' wi_z/wo_z may less or equal than 0, should be 127 | # set to zero. Maybe this step is not necessary, because G is 128 | # already zero. 129 | tmp_zeros = torch.zeros_like(tmp) 130 | static_zero = torch.zeros(1, device=wi.device, dtype=torch.float32) 131 | res = torch.where(torch.le(wi_z, static_zero), tmp_zeros, tmp) 132 | res = torch.where(torch.le(wo_z, static_zero), tmp_zeros, res) 133 | 134 | return res 135 | 136 | @classmethod 137 | def ggx_G1_aniso( 138 | cls, 139 | v: torch.Tensor, 140 | ax: torch.Tensor, 141 | ay: torch.Tensor, 142 | vz: torch.Tensor 143 | ) -> torch.Tensor: 144 | """ 145 | If vz <= 0, return 0 146 | 147 | Args: 148 | v: shape (batch, lightnum, 3) 149 | ax: shape (batch, 1) 150 | ay: shape (batch, 1) 151 | vz: shape (batch, lightnum, 1) 152 | 153 | Returns: 154 | G1: shape (batch, lightnum, 1) 155 | """ 156 | axayaz = torch.cat([ax, ay, torch.ones_like(ax)], dim=1) # (batch, 3) 157 | vv = v * torch.unsqueeze(axayaz, dim=1) # (batch, lightnum, 3) 158 | G1 = 2.0 * vz / (vz + torch.norm(vv, dim=2, keepdim=True) + 1e-6) 159 | 160 | # If vz < 0, G1 will be zero. 161 | G1 = torch.where( 162 | torch.le(vz, torch.zeros_like(vz)), 163 | torch.zeros_like(vz), 164 | G1 165 | ) 166 | return G1 -------------------------------------------------------------------------------- /GGXBRDF/onb.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | 4 | 5 | class ONB(object): 6 | 7 | def __init__(self, batch_size : int): 8 | """ 9 | uvw - xyz - tbn 10 | """ 11 | self.batch_size = batch_size 12 | self.axis = torch.zeros(batch_size, 3, 3) # (batch, 3, 3) 13 | self.axis[:, 0, 0] = 1 14 | self.axis[:, 1, 1] = 1 15 | self.axis[:, 2, 2] = 1 16 | 17 | def u(self) -> torch.Tensor: 18 | """ 19 | Returns: 20 | u axis of shape (batch, 3) 21 | """ 22 | return self.axis[:, 0, :] 23 | 24 | def v(self) -> torch.Tensor: 25 | return self.axis[:, 1, :] 26 | 27 | def w(self) -> torch.Tensor: 28 | return self.axis[:, 2, :] 29 | 30 | def inverse_transform(self, p : torch.Tensor) -> torch.Tensor: 31 | """ 32 | Convert local coordinate(in onb) back to global 33 | coordinate(onb in). 34 | 35 | Args: 36 | p: local coordinates of shape (batch, 3) or (batch, N, 3) 37 | 38 | Returns: 39 | global coordinates of shape (batch, 3) or (batch, N, 3) 40 | """ 41 | assert(self.batch_size == p.size(0)) 42 | assert(len(p.size()) in [2, 3]) 43 | if len(p.size()) == 2: 44 | return p[:, 0:1] * self.u() + p[:, 1:2] * self.v() + p[:, 2:3] * self.w() 45 | elif len(p.size()) == 3: 46 | u = self.u().unsqueeze(1) 47 | v = self.v().unsqueeze(1) 48 | w = self.w().unsqueeze(1) 49 | return p[:, :, [0]] * u + p[:, :, [1]] * v + p[:, :, [2]] * w 50 | 51 | def transform(self, p : torch.Tensor) -> torch.Tensor: 52 | """ 53 | Convert global coordinate(onb in) to local 54 | coordinate(in onb). 55 | 56 | Args: 57 | p: global coordinates of shape (batch, 3) or (batch, lightnum, 3) 58 | 59 | Returns: 60 | local coordinates of shape (batch, 3) or (batch, lightnum, 3) 61 | """ 62 | assert(self.batch_size == p.size(0)) 63 | assert(len(p.size()) in [2, 3]) 64 | if len(p.size()) == 2: 65 | x = torch.sum(p * self.u(), dim=1, keepdim=True) 66 | y = torch.sum(p * self.v(), dim=1, keepdim=True) 67 | z = torch.sum(p * self.w(), dim=1, keepdim=True) 68 | return torch.cat([x, y, z], dim=1) 69 | elif len(p.size()) == 3: 70 | lightnum = p.size(1) 71 | u = self.u().unsqueeze(1).repeat(1, lightnum, 1) 72 | v = self.v().unsqueeze(1).repeat(1, lightnum, 1) 73 | w = self.w().unsqueeze(1).repeat(1, lightnum, 1) 74 | x = torch.sum(p * u, dim=2, keepdim=True) 75 | y = torch.sum(p * v, dim=2, keepdim=True) 76 | z = torch.sum(p * w, dim=2, keepdim=True) 77 | return torch.cat([x, y, z], dim=2) 78 | 79 | def build_from_ntb( 80 | self, 81 | n : torch.Tensor, 82 | t : torch.Tensor, 83 | b : torch.Tensor, 84 | ) -> None: 85 | """ 86 | Args: 87 | n, t, b: The local frame, of shape (batch, 3) 88 | """ 89 | batch_size = n.size(0) 90 | self.axis = torch.zeros((batch_size, 3, 3)).to(n.device) 91 | self.axis[:, 2, :] = F.normalize(n, dim=1) 92 | self.axis[:, 1, :] = F.normalize(b, dim=1) 93 | self.axis[:, 0, :] = F.normalize(t, dim=1) 94 | 95 | def build_from_w(self, normal : torch.Tensor) -> None: 96 | """ 97 | Build the local frame based on the normal. 98 | 99 | Args: 100 | normal: The normal coordinates of shape (batch, 3) 101 | """ 102 | assert(self.batch_size == normal.size(0)) 103 | device = normal.device 104 | n = F.normalize(normal, dim=1) 105 | nz = n[:, [2]] 106 | batch_size = n.shape[0] 107 | 108 | constant_001 = torch.zeros_like(normal).to(device) 109 | constant_001[:, 2] = 1.0 110 | constant_100 = torch.zeros_like(normal).to(device) 111 | constant_100[:, 0] = 1.0 112 | 113 | nz_notequal_1 = torch.gt(torch.abs(nz - 1.0), 1e-6) 114 | nz_notequal_m1 = torch.gt(torch.abs(nz + 1.0), 1e-6) 115 | 116 | 117 | t = torch.where(nz_notequal_1 & nz_notequal_m1, constant_001, constant_100) 118 | # Optix version 119 | # b = F.normalize(torch.cross(normal, t), dim=1) 120 | # t = torch.cross(b, normal) 121 | # Original pytorch version 122 | t = F.normalize(torch.cross(t, normal), dim=1) 123 | b = torch.cross(n, t) 124 | 125 | self.axis = torch.zeros((batch_size, 3, 3)).to(device) 126 | self.axis[:, 2, :] = n 127 | self.axis[:, 1, :] = b 128 | self.axis[:, 0, :] = t 129 | 130 | def rotate_frame(self, theta : torch.Tensor) -> None: 131 | """ 132 | Rotate local frame along the normal axis 133 | 134 | Args: 135 | theta: the degrees of counterclockwise rotation, of shape (batch, 1) 136 | """ 137 | assert(self.batch_size == theta.size(0)) 138 | n = self.w() 139 | t = self.u() 140 | b = self.v() 141 | 142 | t = F.normalize(t * torch.cos(theta) + b * torch.sin(theta), dim=1) 143 | b = F.normalize(torch.cross(n, t), dim=1) 144 | self.axis = torch.zeros((self.batch_size, 3, 3)).to(theta.device) 145 | self.axis[:, 0, :] = t 146 | self.axis[:, 1, :] = b 147 | self.axis[:, 2, :] = n 148 | 149 | def _back_hemi_octa_map(self, n_2d : torch.Tensor) -> torch.Tensor: 150 | """ 151 | The original normal is (0, 0, 1), we should use this method to 152 | perturb the original normal to get a new normal and then build 153 | a new local frame based on the new normal. 154 | 155 | Args: 156 | n_2d: shape (batch, 2) 157 | 158 | Returns: 159 | local_n: shape (batch, 3), which is define in geometry local 160 | frame. 161 | """ 162 | p = (n_2d - 0.5) * 2.0 163 | resultx = (p[:, [0]] + p[:, [1]]) * 0.5 164 | resulty = (p[:, [1]] - p[:, [0]]) * 0.5 165 | resultz = 1.0 - torch.abs(resultx) - torch.abs(resulty) 166 | result = torch.cat([resultx, resulty, resultz], dim=1) 167 | return F.normalize(result, dim=1) 168 | 169 | def hemi_octa_map(self, dir : torch.Tensor) -> torch.Tensor: 170 | """ 171 | Args: 172 | dir: shape (batch, 3) 173 | 174 | Returns: 175 | n2d: shape (batch, 2), which is define in circle coordinate 176 | """ 177 | high_dim = False 178 | batch_size = dir.shape[0] 179 | 180 | if len(dir.shape) > 2: 181 | high_dim = True 182 | dir = dir.reshape(-1, 3) 183 | 184 | p = dir/torch.sum(torch.abs(dir), dim=1, keepdim=True) # (batch,3) 185 | n_2d = torch.cat([p[:,[0]] - p[:,[1]],p[:,[0]] + p[:, [1]]],dim=1) * 0.5 + 0.5 186 | if high_dim: 187 | n_2d = n_2d.reshape(batch_size, -1, 2) 188 | return n_2d 189 | 190 | def build_from_n2d(self, n_2d : torch.Tensor, theta : torch.Tensor) -> None: 191 | """ 192 | Args: 193 | n_2d: tensor of shape (batch, 2). the param defines how 194 | to perturb local normal. 195 | theta: the degrees of rotation of tangent. 196 | """ 197 | assert(self.batch_size == n_2d.size(0)) 198 | 199 | local_n = self._back_hemi_octa_map(n_2d) 200 | self.build_from_w(local_n) 201 | self.rotate_frame(theta) 202 | -------------------------------------------------------------------------------- /assets/sorted_teaser_low_res.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/OpenSVBRDF/OpenSVBRDF_source_code/27d60e6e95bc66f65d1780b114f2aa228992e7f4/assets/sorted_teaser_low_res.png -------------------------------------------------------------------------------- /assets/sorted_teaser_new.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/OpenSVBRDF/OpenSVBRDF_source_code/27d60e6e95bc66f65d1780b114f2aa228992e7f4/assets/sorted_teaser_new.jpg -------------------------------------------------------------------------------- /data_processing/camera_related/camera_config.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | 4 | class Camera(): 5 | """ 6 | Camera encapsulates operations related to camera calibration and undistortion. 7 | 8 | Functions: 9 | project: Project the spatial coordinates onto the image plane. 10 | undistort_2steps: Undistort the source image using cv2.initUndistortRectifyMap() and cv2.remap(). 11 | get_cam_pos: Retrieve the camera position. 12 | get_trans_matrix: Retrieve the rotation and translation matrices. 13 | get_camera_matrix: Retrieve the camera matrix. 14 | get_distortion_matrix: Retrieve the distortion matrix. 15 | get_height: Retrieve the image height. 16 | get_width: Retrieve the image width. 17 | 18 | """ 19 | def __init__(self,intrinsic_file,extrinsic_file): 20 | # Load intrinsic parameters (camera matrix and distortion coefficients) 21 | cv_file = cv2.FileStorage(intrinsic_file, cv2.FILE_STORAGE_READ) 22 | self.A = cv_file.getNode("camera_matrix").mat() 23 | self.D = cv_file.getNode("distortion_coefficients").mat() 24 | cv_file.release() 25 | 26 | # Load extrinsic parameters (rotation vector and translation vector) 27 | cv_file = cv2.FileStorage(extrinsic_file, cv2.FILE_STORAGE_READ) 28 | self.rvec = cv_file.getNode("rvec").mat() 29 | self.tvec = cv_file.getNode("tvec").mat() 30 | self.rvec = np.stack(self.rvec,axis=0) 31 | self.tvec = np.stack(self.tvec,axis=0) 32 | cv_file.release() 33 | 34 | # Convert rotation vector to rotation matrix 35 | self.R = cv2.Rodrigues(self.rvec)[0] 36 | 37 | self.T = self.tvec 38 | 39 | # Compute camera position in the world coordinate system 40 | self.cam_pos = np.matmul(-np.linalg.inv(self.R),self.T) 41 | 42 | # Prepare projection matrix 43 | self.matP = np.zeros([4,4],np.float32) 44 | self.matP[:3,:3] = self.R 45 | self.matP[:3,[3]] = self.T 46 | 47 | # Image dimensions 48 | self.width = 5328 49 | self.height = 4608 50 | 51 | self.ncm, _ = cv2.getOptimalNewCameraMatrix(self.A, self.D, (self.width, self.height), 1, (self.width, self.height), 0) 52 | 53 | # Initialize undistortion and rectification map 54 | self.map1, self.map2 = cv2.initUndistortRectifyMap(self.A, self.D, None, self.ncm, 55 | (self.width, self.height), cv2.CV_16SC2) 56 | 57 | def get_cam_pos(self): 58 | """Retrieve the camera position.""" 59 | return self.cam_pos 60 | 61 | def get_trans_matrix(self): 62 | """Retrieve the rotation and translation matrices.""" 63 | return self.R, self.T 64 | 65 | def get_camera_matrix(self): 66 | """Retrieve the camera matrix.""" 67 | return self.A 68 | 69 | def get_distortion_matrix(self): 70 | """Retrieve the distortion matrix.""" 71 | return self.D 72 | 73 | def project(self, pos): 74 | """ 75 | Project spatial coordinates onto the image plane. 76 | 77 | Parameters: 78 | pos (np.ndarray): Spatial coordinates. 79 | 80 | Returns: 81 | np.ndarray: Image coordinates. 82 | """ 83 | cam_coord_pos = np.matmul(self.R, pos)+self.T 84 | cam_coord_pos /= cam_coord_pos[[2],:] 85 | 86 | cam_coord = np.matmul(self.A, cam_coord_pos).T 87 | 88 | return cam_coord[:,:2] 89 | 90 | def undistort_2steps(self,src): 91 | """ 92 | Undistort the source image using precomputed maps. 93 | 94 | Parameters: 95 | src (np.ndarray): Source image. 96 | 97 | Returns: 98 | np.ndarray: Undistorted image. 99 | """ 100 | dst = cv2.remap(src, self.map1, self.map2, cv2.INTER_LINEAR) 101 | return dst 102 | 103 | def get_height(self): 104 | """Retrieve the image height.""" 105 | return self.height 106 | 107 | def get_width(self): 108 | """Retrieve the image width.""" 109 | return self.width 110 | -------------------------------------------------------------------------------- /data_processing/extract_measurements/run.sh: -------------------------------------------------------------------------------- 1 | for variable in 18 2 | do 3 | class_name=paper 4 | formatted_variable=$(printf "%04d" $variable) 5 | data_root=../../database_data/"$class_name""$formatted_variable"/ 6 | 7 | dir_name=raw_images 8 | 9 | images_number=129 10 | cam_num=2 11 | main_cam_id=0 12 | texture_resolution=1024 13 | 14 | save_root=../../database_data/$class_name"$formatted_variable"/output/texture_"$texture_resolution"/ 15 | 16 | model_path=../../database_model/ 17 | config_dir=../device_configuration/ 18 | 19 | need_undistort=true 20 | color_check=true 21 | need_scale=true 22 | need_warp=true 23 | phto_validation=false 24 | down_size=2 25 | 26 | lighting_pattern_num=64 27 | line_pattern_num=64 28 | 29 | python extract.py $data_root $dir_name $save_root $images_number $cam_num $main_cam_id $config_dir $model_path $texture_resolution $lighting_pattern_num --line_pattern_num $line_pattern_num --down_size $down_size --need_undistort $need_undistort --color_check $color_check --need_scale $need_scale --need_warp $need_warp 30 | 31 | done 32 | -------------------------------------------------------------------------------- /data_processing/finetune/AUTO_lumitexel_net_inference.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | import numpy as np 5 | from collections import OrderedDict 6 | import math 7 | 8 | class LumitexelNet(nn.Module): 9 | ''' 10 | LumitexelNet defines a neural network model for lumitexel. 11 | 12 | Functions: 13 | color_latent_part: submodel for processing color information of a latent. 14 | 15 | latent_part: submodel for latent part. 16 | 17 | lumi_net: submodel for lumitextel. 18 | 19 | ''' 20 | def __init__(self,args): 21 | super(LumitexelNet,self).__init__() 22 | 23 | self.input_length = args["lumitexel_length"] 24 | 25 | self.shape_latent_len = args["shape_latent_len"] 26 | self.color_latent_len = args["color_latent_len"] 27 | self.latent_len = self.shape_latent_len + self.color_latent_len 28 | # construct model 29 | self.latent_part_model = self.latent_part(self.input_length) 30 | self.color_latent_part_model = self.color_latent_part(self.input_length) 31 | self.lumi_net_model = self.lumi_net(self.latent_len) 32 | 33 | def color_latent_part(self,input_size,name_prefix="Color_Latent_"): 34 | 35 | 36 | layer_stack = OrderedDict() 37 | 38 | layer_count = 0 39 | 40 | output_size=input_size // 4 41 | layer_stack[name_prefix+"Linear_{}".format(layer_count)] = nn.Linear(input_size,output_size) 42 | layer_stack[name_prefix+"LeakyRelu_{}".format(layer_count)] = nn.LeakyReLU() 43 | layer_count+=1 44 | input_size = output_size 45 | 46 | 47 | output_size=1024 48 | layer_stack[name_prefix+"Linear_{}".format(layer_count)] = nn.Linear(input_size,output_size) 49 | layer_stack[name_prefix+"LeakyRelu_{}".format(layer_count)] = nn.LeakyReLU() 50 | layer_count+=1 51 | input_size = output_size 52 | 53 | 54 | output_size=256 55 | layer_stack[name_prefix+"BN_{}".format(layer_count)] = nn.BatchNorm1d(input_size) 56 | layer_stack[name_prefix+"Linear_{}".format(layer_count)] = nn.Linear(input_size,output_size) 57 | layer_stack[name_prefix+"LeakyRelu_{}".format(layer_count)] = nn.LeakyReLU() 58 | layer_count+=1 59 | input_size = output_size 60 | 61 | output_size=64 62 | layer_stack[name_prefix+"BN_{}".format(layer_count)] = nn.BatchNorm1d(input_size) 63 | layer_stack[name_prefix+"Linear_{}".format(layer_count)] = nn.Linear(input_size,output_size) 64 | layer_stack[name_prefix+"LeakyRelu_{}".format(layer_count)] = nn.LeakyReLU() 65 | layer_count+=1 66 | input_size = output_size 67 | 68 | output_size=self.color_latent_len 69 | layer_stack[name_prefix+"Linear_{}".format(layer_count)] = nn.Linear(input_size,output_size) 70 | layer_count+=1 71 | input_size = output_size 72 | 73 | layer_stack = nn.Sequential(layer_stack) 74 | 75 | return layer_stack 76 | 77 | def latent_part(self,input_size,name_prefix="Latent_"): 78 | 79 | layer_stack = OrderedDict() 80 | 81 | layer_count = 0 82 | 83 | output_size=input_size // 2 84 | layer_stack[name_prefix+"Linear_{}".format(layer_count)] = nn.Linear(input_size,output_size) 85 | layer_stack[name_prefix+"LeakyRelu_{}".format(layer_count)] = nn.LeakyReLU() 86 | layer_count+=1 87 | input_size = output_size 88 | 89 | output_size=4096 90 | layer_stack[name_prefix+"Linear_{}".format(layer_count)] = nn.Linear(input_size,output_size) 91 | layer_stack[name_prefix+"LeakyRelu_{}".format(layer_count)] = nn.LeakyReLU() 92 | layer_count+=1 93 | input_size = output_size 94 | 95 | output_size=2048 96 | layer_stack[name_prefix+"BN_{}".format(layer_count)] = nn.BatchNorm1d(input_size) 97 | layer_stack[name_prefix+"Linear_{}".format(layer_count)] = nn.Linear(input_size,output_size) 98 | layer_stack[name_prefix+"LeakyRelu_{}".format(layer_count)] = nn.LeakyReLU() 99 | layer_count+=1 100 | input_size = output_size 101 | 102 | output_size=1024 103 | layer_stack[name_prefix+"BN_{}".format(layer_count)] = nn.BatchNorm1d(input_size) 104 | layer_stack[name_prefix+"Linear_{}".format(layer_count)] = nn.Linear(input_size,output_size) 105 | layer_stack[name_prefix+"LeakyRelu_{}".format(layer_count)] = nn.LeakyReLU() 106 | layer_count+=1 107 | input_size = output_size 108 | 109 | output_size=512 110 | layer_stack[name_prefix+"BN_{}".format(layer_count)] = nn.BatchNorm1d(input_size) 111 | layer_stack[name_prefix+"Linear_{}".format(layer_count)] = nn.Linear(input_size,output_size) 112 | layer_stack[name_prefix+"LeakyRelu_{}".format(layer_count)] = nn.LeakyReLU() 113 | layer_count+=1 114 | input_size = output_size 115 | 116 | output_size=256 117 | layer_stack[name_prefix+"BN_{}".format(layer_count)] = nn.BatchNorm1d(input_size) 118 | layer_stack[name_prefix+"Linear_{}".format(layer_count)] = nn.Linear(input_size,output_size) 119 | layer_stack[name_prefix+"LeakyRelu_{}".format(layer_count)] = nn.LeakyReLU() 120 | layer_count+=1 121 | input_size = output_size 122 | 123 | output_size=256 124 | layer_stack[name_prefix+"BN_{}".format(layer_count)] = nn.BatchNorm1d(input_size) 125 | layer_stack[name_prefix+"Linear_{}".format(layer_count)] = nn.Linear(input_size,output_size) 126 | layer_stack[name_prefix+"LeakyRelu_{}".format(layer_count)] = nn.LeakyReLU() 127 | layer_count+=1 128 | input_size = output_size 129 | 130 | output_size=self.shape_latent_len 131 | layer_stack[name_prefix+"Linear_{}".format(layer_count)] = nn.Linear(input_size,output_size) 132 | layer_count+=1 133 | input_size = output_size 134 | 135 | layer_stack = nn.Sequential(layer_stack) 136 | 137 | return layer_stack 138 | 139 | def lumi_net(self,input_size,name_prefix = "Lumi_"): 140 | 141 | layer_stack = OrderedDict() 142 | 143 | layer_count = 0 144 | 145 | output_size=128 146 | layer_stack[name_prefix+"Linear_{}".format(layer_count)] = nn.Linear(input_size,output_size) 147 | layer_stack[name_prefix+"LeakyRelu_{}".format(layer_count)] = nn.LeakyReLU() 148 | layer_count+=1 149 | input_size = output_size 150 | 151 | output_size=256 152 | layer_stack[name_prefix+"Linear_{}".format(layer_count)] = nn.Linear(input_size,output_size) 153 | layer_stack[name_prefix+"LeakyRelu_{}".format(layer_count)] = nn.LeakyReLU() 154 | layer_count+=1 155 | input_size = output_size 156 | 157 | output_size=256 158 | layer_stack[name_prefix+"Linear_{}".format(layer_count)] = nn.Linear(input_size,output_size) 159 | layer_stack[name_prefix+"LeakyRelu_{}".format(layer_count)] = nn.LeakyReLU() 160 | layer_count+=1 161 | input_size = output_size 162 | 163 | 164 | output_size=512 165 | layer_stack[name_prefix+"Linear_{}".format(layer_count)] = nn.Linear(input_size,output_size) 166 | layer_stack[name_prefix+"LeakyRelu_{}".format(layer_count)] = nn.LeakyReLU() 167 | layer_count+=1 168 | input_size = output_size 169 | 170 | output_size=1024 171 | layer_stack[name_prefix+"Linear_{}".format(layer_count)] = nn.Linear(input_size,output_size) 172 | layer_stack[name_prefix+"LeakyRelu_{}".format(layer_count)] = nn.LeakyReLU() 173 | layer_count+=1 174 | input_size = output_size 175 | 176 | output_size=2048 177 | layer_stack[name_prefix+"Linear_{}".format(layer_count)] = nn.Linear(input_size,output_size) 178 | layer_stack[name_prefix+"LeakyRelu_{}".format(layer_count)] = nn.LeakyReLU() 179 | layer_count+=1 180 | input_size = output_size 181 | 182 | output_size=2048 183 | layer_stack[name_prefix+"Linear_{}".format(layer_count)] = nn.Linear(input_size,output_size) 184 | layer_stack[name_prefix+"LeakyRelu_{}".format(layer_count)] = nn.LeakyReLU() 185 | layer_count+=1 186 | input_size = output_size 187 | 188 | output_size=4096 189 | layer_stack[name_prefix+"Linear_{}".format(layer_count)] = nn.Linear(input_size,output_size) 190 | layer_stack[name_prefix+"LeakyRelu_{}".format(layer_count)] = nn.LeakyReLU() 191 | layer_count+=1 192 | input_size = output_size 193 | 194 | output_size=4096 195 | layer_stack[name_prefix+"Linear_{}".format(layer_count)] = nn.Linear(input_size,output_size) 196 | layer_stack[name_prefix+"LeakyRelu_{}".format(layer_count)] = nn.LeakyReLU() 197 | layer_count+=1 198 | input_size = output_size 199 | 200 | output_size=8192 201 | layer_stack[name_prefix+"Linear_{}".format(layer_count)] = nn.Linear(input_size,output_size) 202 | layer_stack[name_prefix+"LeakyRelu_{}".format(layer_count)] = nn.LeakyReLU() 203 | layer_count+=1 204 | input_size = output_size 205 | 206 | output_size=self.input_length 207 | layer_stack[name_prefix+"Linear_{}".format(layer_count)] = nn.Linear(input_size,output_size) 208 | 209 | layer_stack = nn.Sequential(layer_stack) 210 | 211 | return layer_stack 212 | 213 | 214 | def forward(self,net_input,input_is_latent=False,return_latent_directly=False): 215 | batch_size = net_input.shape[0] 216 | net_input = net_input.reshape([batch_size,-1]) 217 | 218 | if input_is_latent is True: 219 | latent = net_input 220 | else: 221 | shape_latent = self.latent_part_model(net_input) 222 | color_latent = self.color_latent_part_model(net_input) 223 | latent = torch.cat([color_latent,shape_latent],dim=-1) 224 | 225 | if return_latent_directly: 226 | return latent 227 | 228 | 229 | nn_lumi = self.lumi_net_model(latent) 230 | 231 | nn_lumi = torch.exp(nn_lumi)-1.0 232 | 233 | return latent,nn_lumi -------------------------------------------------------------------------------- /data_processing/finetune/AUTO_planar_scanner_net_inference.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from AUTO_lumitexel_net_inference import LumitexelNet 3 | 4 | 5 | class PlanarScannerNet(nn.Module): 6 | def __init__(self, args): 7 | super(PlanarScannerNet, self).__init__() 8 | 9 | self.lumitexel_net = LumitexelNet(args) 10 | 11 | 12 | def forward(self, input, input_is_latent=True): 13 | nn_latent,nn_lumi = self.lumitexel_net(input, input_is_latent=input_is_latent) 14 | 15 | return nn_latent,nn_lumi 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | -------------------------------------------------------------------------------- /data_processing/finetune/choose_branch_net.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | import numpy as np 5 | from collections import OrderedDict 6 | import math 7 | 8 | class ChooseBranchNet(nn.Module): 9 | ''' 10 | ChooseBranchNet defines a neural network model for selecting branches. 11 | 12 | functions: 13 | branch_decoder: build a branch decoder model, including a series of convolutional layers, batch normalization layers and activation function layers. 14 | It returns a sequential container containing all the layers and initializing the weight and bias of the convolutional layer. 15 | 16 | forward: use forward propagation to obtain predicted results. 17 | It returns reshaped predicted results. 18 | 19 | by Leyao 20 | ''' 21 | def __init__(self,args): 22 | super(ChooseBranchNet,self).__init__() 23 | 24 | self.classifier_num = args["classifier_num"] 25 | 26 | self.layers = args["layers"] 27 | self.input_length = args["lighting_pattern_num"] * (self.layers-1) 28 | self.layer_width = [64,64,32,16,8,2] 29 | self.layer_width = [value * (self.layers-1) for value in self.layer_width] 30 | 31 | # construct model 32 | 33 | self.choose_branch_net_model = self.branch_decoder() 34 | 35 | def branch_decoder(self,name_prefix = "choose_branch_net_"): 36 | layer_stack = OrderedDict() 37 | 38 | layer_count = 0 39 | input_size = self.input_length 40 | 41 | for which_layer in self.layer_width[:-1]: 42 | output_size = which_layer 43 | layer_stack[name_prefix+"Conv_{}".format(layer_count)] = torch.nn.Conv1d(input_size, output_size, kernel_size=1,groups=self.layers-1) 44 | layer_stack[name_prefix+"BN_{}".format(layer_count)] = nn.BatchNorm1d(output_size) 45 | layer_stack[name_prefix+"LeakyRelu_{}".format(layer_count)] = nn.LeakyReLU() 46 | 47 | layer_count+=1 48 | input_size = output_size 49 | 50 | output_size = self.layer_width[-1] 51 | layer_stack[name_prefix+"Conv_{}".format(layer_count)] = torch.nn.Conv1d(input_size, output_size, kernel_size=1,groups=self.layers-1) 52 | 53 | layer_stack = nn.Sequential(layer_stack) 54 | 55 | for m in layer_stack: 56 | if isinstance(m, torch.nn.Conv1d): 57 | m.weight.data.normal_(mean = 0, std = 0.1) 58 | m.bias.data.fill_(0.0) 59 | 60 | return layer_stack 61 | 62 | def forward(self,x_n): 63 | pred = self.choose_branch_net_model(x_n) 64 | pred = pred.reshape([-1,self.layers-1,2]) 65 | pred = torch.softmax(pred,dim=-1) 66 | return pred 67 | -------------------------------------------------------------------------------- /data_processing/finetune/finetune_pass1.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the experimental code for paper ``Ma et al., OpenSVBRDF: A Database of Measured Spatially-Varying Reflectance, SIGGRAPH Asia 2023``. 3 | This script is suboptimal and experimental. 4 | There may be redundant lines and functionalities. 5 | 6 | Xiaohe Ma, 2024/02 7 | ''' 8 | 9 | ''' 10 | This script implements the first step of the per-pixel fine-tuning process. 11 | 12 | ''' 13 | 14 | import numpy as np 15 | import argparse 16 | import torch 17 | import os 18 | from datetime import datetime 19 | from torch.autograd import Variable 20 | import AUTO_planar_scanner_net_inference 21 | 22 | 23 | if __name__ == "__main__": 24 | parser = argparse.ArgumentParser() 25 | parser.add_argument("--data_for_server_root",default="data_for_server/") 26 | parser.add_argument("--lighting_pattern_num",type=int,default=64) 27 | parser.add_argument("--finetune_use_num",type=int,default=64) 28 | parser.add_argument("--batchsize",type=int,default=100) 29 | parser.add_argument('--thread_ids', nargs='+', type=int,default=[5]) 30 | parser.add_argument('--gpu_id', type=int,default=0) 31 | parser.add_argument("--need_dump",action="store_true") 32 | parser.add_argument("--total_thread_num",type=int,default=24) 33 | parser.add_argument("--tex_resolution",type=int,default=512) 34 | parser.add_argument("--cam_num",type=int,default=2) 35 | parser.add_argument("--main_cam_id",type=int,default=0) 36 | parser.add_argument("--model_file",type=str,default="../../model/model_state_450000.pkl") 37 | parser.add_argument("--pattern_file",type=str,default="../../model/opt_W.bin") 38 | 39 | parser.add_argument("--shape_latent_len",type=int,default=2) 40 | parser.add_argument("--color_latent_len",type=int,default=2) 41 | parser.add_argument("--save_lumi",action="store_true") 42 | 43 | args = parser.parse_args() 44 | 45 | torch.set_printoptions(precision=3) 46 | 47 | compute_device = torch.device(f"cuda:{args.gpu_id}") 48 | 49 | train_configs = {} 50 | train_configs["training_device"] = args.gpu_id 51 | train_configs["lumitexel_length"] = 64*64*3 52 | train_configs["shape_latent_len"] = args.shape_latent_len 53 | train_configs["color_latent_len"] = args.color_latent_len 54 | 55 | latent_len = args.shape_latent_len + args.color_latent_len 56 | 57 | pretrained_dict = torch.load(args.model_file, map_location=compute_device) 58 | inference_net = AUTO_planar_scanner_net_inference.PlanarScannerNet(train_configs) 59 | 60 | m_len_perview = 3 61 | 62 | something_not_found = False 63 | model_dict = inference_net.state_dict() 64 | for k,_ in model_dict.items(): 65 | if k not in pretrained_dict: 66 | print("not found:", k) 67 | something_not_found = True 68 | if something_not_found: 69 | exit() 70 | 71 | model_dict = inference_net.state_dict() 72 | pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} 73 | model_dict.update(pretrained_dict) 74 | inference_net.load_state_dict(model_dict) 75 | for p in inference_net.parameters(): 76 | p.requires_grad=False 77 | inference_net.to(compute_device) 78 | 79 | inference_net.eval() 80 | 81 | for which_test_thread in args.thread_ids: 82 | data_root = args.data_for_server_root+f"{which_test_thread}/" 83 | print(data_root) 84 | 85 | lighting_patterns_np = np.fromfile(args.pattern_file,np.float32).reshape([args.lighting_pattern_num,-1,3])[:args.finetune_use_num,:,0].reshape([1,args.finetune_use_num,-1]) 86 | 87 | lighting_patterns = torch.from_numpy(lighting_patterns_np).to(compute_device) 88 | 89 | log_path = data_root+"lumi_imgs/" 90 | os.makedirs(log_path,exist_ok=True) 91 | 92 | 93 | measurements = torch.from_numpy(np.fromfile(data_root+f"gt_measurements_{args.tex_resolution}.bin",np.float32)).to(compute_device).reshape((-1,args.lighting_pattern_num,m_len_perview))[:,:args.finetune_use_num,:] 94 | 95 | 96 | pf_nn_latent = np.fromfile(data_root+f"latent_{args.tex_resolution}.bin",np.float32).reshape([-1,latent_len]) 97 | latent_num = pf_nn_latent.shape[0] 98 | 99 | assert measurements.shape[0] == latent_num,"some data are corrupted" 100 | sample_num = measurements.shape[0] 101 | 102 | pf_result_grey = open(data_root+f"pass1_latent_{args.tex_resolution}.bin","wb") 103 | pf_result_nn = open(data_root+f"pass0_latent_{args.tex_resolution}.bin","wb") 104 | 105 | ptr = 0 106 | optimize_step = 500 107 | lr = 0.02 108 | 109 | print("finetune latent...") 110 | 111 | while True: 112 | if ptr % 30000 == 0: 113 | start = datetime.now() 114 | print(f"PASS 1 [{which_test_thread}]/{ptr}/{sample_num} {start}") 115 | 116 | tmp_measurements = measurements[ptr:ptr+args.batchsize] 117 | tmp_measurements_mean = tmp_measurements.mean(dim=-1) 118 | 119 | cur_batchsize = tmp_measurements.shape[0] 120 | if cur_batchsize == 0: 121 | print("break because all done.") 122 | break 123 | 124 | tmp_nn_latent = pf_nn_latent[ptr:ptr+cur_batchsize] 125 | tmp_nn_color_latent = pf_nn_latent[ptr:ptr+cur_batchsize,:args.color_latent_len] 126 | tmp_nn_shape_latent = pf_nn_latent[ptr:ptr+cur_batchsize,args.color_latent_len:] 127 | tmp_nn_latent_3c = np.concatenate([tmp_nn_color_latent, tmp_nn_color_latent,tmp_nn_color_latent,tmp_nn_shape_latent],axis=1) 128 | tmp_nn_latent_3c.astype(np.float32).tofile(pf_result_nn) 129 | 130 | tmp_x_guess = torch.cuda.FloatTensor(tmp_nn_latent,device=compute_device) 131 | tmp_x_guess = Variable(tmp_x_guess,requires_grad=True) 132 | optimizer = torch.optim.Adam([tmp_x_guess,], lr = lr) 133 | 134 | latent_start = datetime.now() 135 | 136 | loss_step = [] 137 | loss_precent = [] 138 | for step in range(optimize_step): 139 | _,tmp_lumi = inference_net(tmp_x_guess,input_is_latent=True) 140 | tmp_lumi = torch.max(torch.zeros_like(tmp_lumi),tmp_lumi) 141 | 142 | tmp_lumi_measurements = torch.sum(lighting_patterns*tmp_lumi.unsqueeze(dim=1),dim=-1).reshape([cur_batchsize,-1]) 143 | 144 | loss = torch.nn.functional.mse_loss(torch.pow(tmp_lumi_measurements,1/2.0), torch.pow(tmp_measurements_mean,1/2.0),reduction='sum') 145 | 146 | optimizer.zero_grad() 147 | loss.backward() 148 | optimizer.step() 149 | 150 | loss_step.append(loss) 151 | if step >= 10: 152 | tmp_loss_percent = loss/loss_step[step-10]*100 153 | loss_precent.append(tmp_loss_percent) 154 | if len(loss_precent) >= 5: 155 | if torch.mean(torch.stack(loss_precent[-5:],dim=0)) > 95.0: 156 | break 157 | 158 | tmp_x_guess = tmp_x_guess.detach() 159 | 160 | batch_collector = [] 161 | 162 | nn_color_lumi = [] 163 | 164 | tmp_shape_latent = tmp_x_guess[:,args.color_latent_len:] 165 | 166 | tmp_color_latent = tmp_x_guess[:,:args.color_latent_len].unsqueeze(dim=1).repeat(1,3,1).reshape([-1,args.color_latent_len*3]) 167 | grey_latent_3channel = torch.cat([tmp_color_latent,tmp_shape_latent],dim=-1) 168 | grey_latent_3channel.cpu().numpy().astype(np.float32).tofile(pf_result_grey) 169 | 170 | 171 | ptr = ptr+cur_batchsize 172 | 173 | pf_result_grey.close() 174 | pf_result_nn.close() 175 | 176 | print("done.") 177 | -------------------------------------------------------------------------------- /data_processing/finetune/finetune_pass2.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the experimental code for paper ``Ma et al., OpenSVBRDF: A Database of Measured Spatially-Varying Reflectance, SIGGRAPH Asia 2023``. 3 | This script is suboptimal and experimental. 4 | There may be redundant lines and functionalities. 5 | 6 | Xiaohe Ma, 2024/02 7 | ''' 8 | 9 | ''' 10 | This script implements the second step of the per-pixel fine-tuning process. 11 | 12 | ''' 13 | 14 | 15 | 16 | import numpy as np 17 | import argparse 18 | import torch 19 | import sys 20 | import os 21 | from datetime import datetime 22 | from torch.autograd import Variable 23 | import AUTO_planar_scanner_net_inference 24 | 25 | 26 | 27 | if __name__ == "__main__": 28 | parser = argparse.ArgumentParser() 29 | parser.add_argument("--data_for_server_root",default="data_for_server/") 30 | parser.add_argument("--lighting_pattern_num",type=int,default="128") 31 | parser.add_argument("--finetune_use_num",type=int,default="128") 32 | parser.add_argument("--batchsize",type=int,default=100) 33 | parser.add_argument('--thread_ids', nargs='+', type=int,default=[5]) 34 | parser.add_argument('--gpu_id', type=int,default=0) 35 | parser.add_argument("--need_dump",action="store_true") 36 | parser.add_argument("--total_thread_num",type=int,default=24) 37 | parser.add_argument("--tex_resolution",type=int,default=512) 38 | parser.add_argument("--cam_num",type=int,default=2) 39 | parser.add_argument("--main_cam_id",type=int,default=0) 40 | parser.add_argument("--model_file",type=str,default="../../model/model_state_450000.pkl") 41 | parser.add_argument("--pattern_file",type=str,default="../../model/opt_W.bin") 42 | parser.add_argument("--shape_latent_len",type=int,default=2) 43 | parser.add_argument("--color_latent_len",type=int,default=2) 44 | parser.add_argument("--save_lumi",action="store_true") 45 | 46 | args = parser.parse_args() 47 | compute_device = torch.device("cuda:{}".format(args.gpu_id)) 48 | 49 | train_configs = {} 50 | train_configs["training_device"] = args.gpu_id 51 | train_configs["lumitexel_length"] = 64*64*3 52 | train_configs["shape_latent_len"] = args.shape_latent_len 53 | train_configs["color_latent_len"] = args.color_latent_len 54 | 55 | all_latent_len = args.shape_latent_len + args.color_latent_len * 3 56 | 57 | pretrained_dict = torch.load(args.model_file, map_location=compute_device) 58 | inference_net = AUTO_planar_scanner_net_inference.PlanarScannerNet(train_configs) 59 | 60 | m_len_perview = 3 61 | 62 | something_not_found = False 63 | model_dict = inference_net.state_dict() 64 | for k,_ in model_dict.items(): 65 | if k not in pretrained_dict: 66 | print("not found:", k) 67 | something_not_found = True 68 | if something_not_found: 69 | exit() 70 | 71 | model_dict = inference_net.state_dict() 72 | pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} 73 | model_dict.update(pretrained_dict) 74 | inference_net.load_state_dict(model_dict) 75 | for p in inference_net.parameters(): 76 | p.requires_grad=False 77 | inference_net.to(compute_device) 78 | 79 | inference_net.eval() 80 | 81 | 82 | for which_test_thread in args.thread_ids: 83 | data_root = args.data_for_server_root+"{}/".format(which_test_thread) 84 | print(data_root) 85 | 86 | lighting_patterns_np = np.fromfile(args.pattern_file,np.float32).reshape([args.lighting_pattern_num,-1,3])[:args.finetune_use_num,:,0].reshape([1,args.finetune_use_num,-1]) 87 | 88 | lighting_patterns = torch.from_numpy(lighting_patterns_np).to(compute_device) 89 | 90 | log_path = data_root+"lumi_imgs/" 91 | os.makedirs(log_path,exist_ok=True) 92 | 93 | 94 | measurements = torch.from_numpy(np.fromfile(data_root+"gt_measurements_{}.bin".format(args.tex_resolution),np.float32)).to(compute_device).reshape((-1,args.lighting_pattern_num,m_len_perview))[:,:args.finetune_use_num,:] 95 | 96 | pf_pass1_latent = np.fromfile(data_root+"pass1_latent_{}.bin".format(args.tex_resolution),np.float32).reshape([-1,all_latent_len]) 97 | latent_num = pf_pass1_latent.shape[0] 98 | pf_pass1_latent = torch.from_numpy(pf_pass1_latent).to(compute_device) 99 | 100 | assert measurements.shape[0] == latent_num,"some data are corrupted" 101 | sample_num = measurements.shape[0] 102 | 103 | pf_result = open(data_root+"pass2_latent_{}.bin".format(args.tex_resolution),"wb") 104 | 105 | ptr = 0 106 | 107 | color_optimize_step = 50 108 | color_lr = 0.05 109 | 110 | 111 | while True: 112 | 113 | if ptr % 30000 == 0: 114 | start = datetime.now() 115 | print(f"PASS 2 [{which_test_thread}]/{ptr}/{sample_num} {start}") 116 | 117 | tmp_measurements = measurements[ptr:ptr+args.batchsize] 118 | 119 | cur_batchsize = tmp_measurements.shape[0] 120 | if cur_batchsize == 0: 121 | print("break because all done.") 122 | break 123 | 124 | 125 | tmp_pass1_color_latent = pf_pass1_latent[ptr:ptr+cur_batchsize,:args.color_latent_len*3] 126 | tmp_pass1_shape_latent = pf_pass1_latent[ptr:ptr+cur_batchsize,args.color_latent_len*3:] 127 | 128 | latent_start = datetime.now() 129 | batch_collector = [] 130 | pass2_color_lumi = [] 131 | 132 | for which_channel in range(3): 133 | tmp_channel_measurements = tmp_measurements[:,:,which_channel].reshape([cur_batchsize,-1]) 134 | 135 | tmp_color_guess = tmp_pass1_color_latent[:,args.color_latent_len*which_channel:args.color_latent_len*(which_channel+1)].clone() 136 | 137 | tmp_color_guess = Variable(tmp_color_guess,requires_grad=True) 138 | 139 | color_optimizer = torch.optim.Adam([tmp_color_guess,], lr = color_lr) 140 | 141 | loss_step = [] 142 | loss_precent = [] 143 | for step in range(color_optimize_step): 144 | tmp_color_shape_guess = torch.cat([tmp_color_guess,tmp_pass1_shape_latent],dim=-1) 145 | _,tmp_channel_lumi = inference_net(tmp_color_shape_guess,input_is_latent=True) 146 | tmp_channel_lumi = torch.max(torch.zeros_like(tmp_channel_lumi),tmp_channel_lumi) 147 | 148 | tmp_lumi_measurements = torch.sum(lighting_patterns*tmp_channel_lumi.unsqueeze(dim=1),dim=-1).reshape([cur_batchsize,-1]) 149 | 150 | color_loss = torch.nn.functional.mse_loss(torch.pow(tmp_lumi_measurements,1/2.0), torch.pow(tmp_channel_measurements,1/2.0),reduction='sum') 151 | 152 | color_optimizer.zero_grad() 153 | color_loss.backward() 154 | color_optimizer.step() 155 | 156 | loss_step.append(color_loss) 157 | if step >= 5: 158 | tmp_loss_percent = color_loss/loss_step[step-5]*100 159 | loss_precent.append(tmp_loss_percent) 160 | if len(loss_precent) >= 5: 161 | if torch.mean(torch.stack(loss_precent[-5:],dim=0)) > 95.0: 162 | break 163 | 164 | tmp_color_guess = tmp_color_guess.detach() 165 | 166 | batch_collector.append(tmp_color_guess) 167 | pass2_color_lumi.append(tmp_channel_lumi) 168 | 169 | batch_collector.append(tmp_pass1_shape_latent) 170 | batch_collector = torch.cat(batch_collector,dim=-1) 171 | pass2_color_lumi = torch.stack(pass2_color_lumi,dim=-1) 172 | 173 | batch_collector.cpu().numpy().astype(np.float32).tofile(pf_result) 174 | 175 | ptr = ptr+cur_batchsize 176 | 177 | pf_result.close() 178 | 179 | print("done.") 180 | 181 | 182 | 183 | 184 | 185 | 186 | 187 | 188 | 189 | 190 | 191 | 192 | 193 | 194 | 195 | 196 | 197 | 198 | 199 | 200 | 201 | 202 | -------------------------------------------------------------------------------- /data_processing/finetune/finetune_pass3.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the experimental code for paper ``Ma et al., OpenSVBRDF: A Database of Measured Spatially-Varying Reflectance, SIGGRAPH Asia 2023``. 3 | This script is suboptimal and experimental. 4 | There may be redundant lines and functionalities. 5 | 6 | Xiaohe Ma, 2024/02 7 | ''' 8 | 9 | ''' 10 | This script implements the third step of the per-pixel fine-tuning process. 11 | 12 | ''' 13 | import numpy as np 14 | import argparse 15 | import torch 16 | import os 17 | from datetime import datetime 18 | from torch.autograd import Variable 19 | 20 | import AUTO_planar_scanner_net_inference 21 | 22 | 23 | if __name__ == "__main__": 24 | parser = argparse.ArgumentParser() 25 | parser.add_argument("--data_for_server_root",default="data_for_server/") 26 | parser.add_argument("--lighting_pattern_num",type=int,default="128") 27 | parser.add_argument("--finetune_use_num",type=int,default="128") 28 | parser.add_argument("--batchsize",type=int,default=100) 29 | parser.add_argument('--thread_ids', nargs='+', type=int,default=[5]) 30 | parser.add_argument('--gpu_id', type=int,default=0) 31 | parser.add_argument("--need_dump",action="store_true") 32 | parser.add_argument("--total_thread_num",type=int,default=24) 33 | parser.add_argument("--tex_resolution",type=int,default=512) 34 | parser.add_argument("--cam_num",type=int,default=2) 35 | parser.add_argument("--main_cam_id",type=int,default=0) 36 | parser.add_argument("--model_file",type=str,default="../../model/model_state_450000.pkl") 37 | parser.add_argument("--pattern_file",type=str,default="../../model/opt_W.bin.pkl") 38 | parser.add_argument("--shape_latent_len",type=int,default=2) 39 | parser.add_argument("--color_latent_len",type=int,default=2) 40 | parser.add_argument("--save_lumi",action="store_true") 41 | parser.add_argument("--if_continue",type=int,default=0) 42 | 43 | args = parser.parse_args() 44 | compute_device = torch.device("cuda:{}".format(args.gpu_id)) 45 | 46 | 47 | train_configs = {} 48 | train_configs["training_device"] = args.gpu_id 49 | train_configs["lumitexel_length"] = 64*64*3 50 | train_configs["shape_latent_len"] = args.shape_latent_len 51 | train_configs["color_latent_len"] = args.color_latent_len 52 | 53 | latent_len = args.shape_latent_len + args.color_latent_len 54 | 55 | all_latent_len = 3*args.color_latent_len + args.shape_latent_len 56 | 57 | pretrained_dict = torch.load(args.model_file, map_location=compute_device) 58 | inference_net = AUTO_planar_scanner_net_inference.PlanarScannerNet(train_configs) 59 | 60 | m_len_perview = 3 61 | 62 | something_not_found = False 63 | model_dict = inference_net.state_dict() 64 | for k,_ in model_dict.items(): 65 | if k not in pretrained_dict: 66 | print("not found:", k) 67 | something_not_found = True 68 | if something_not_found: 69 | exit() 70 | 71 | model_dict = inference_net.state_dict() 72 | pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} 73 | model_dict.update(pretrained_dict) 74 | inference_net.load_state_dict(model_dict) 75 | for p in inference_net.parameters(): 76 | p.requires_grad=False 77 | inference_net.to(compute_device) 78 | 79 | inference_net.eval() 80 | 81 | for which_test_thread in args.thread_ids: 82 | data_root = args.data_for_server_root+"{}/".format(which_test_thread) 83 | print(data_root) 84 | 85 | lighting_patterns_np = np.fromfile(args.pattern_file,np.float32).reshape([args.lighting_pattern_num,-1,3])[:args.finetune_use_num,:,:].reshape([1,args.finetune_use_num,-1,3]) 86 | 87 | lighting_patterns = torch.from_numpy(lighting_patterns_np).to(compute_device) 88 | 89 | log_path = data_root+"lumi_imgs/" 90 | os.makedirs(log_path,exist_ok=True) 91 | 92 | measurements = torch.from_numpy(np.fromfile(data_root+f"gt_measurements_{args.tex_resolution}.bin",np.float32)).to(compute_device).reshape((-1,args.lighting_pattern_num,m_len_perview))[:,:args.finetune_use_num,:] 93 | 94 | if os.path.isfile(data_root+f"pass3_latent_{args.tex_resolution}.bin") and args.if_continue: 95 | finetune_latent = np.fromfile(data_root+f"pass3_latent_{args.tex_resolution}.bin",np.float32).reshape([-1,all_latent_len]) 96 | print("load from 3") 97 | else: 98 | finetune_latent = np.fromfile(data_root+f"pass2_latent_{args.tex_resolution}.bin",np.float32).reshape([-1,all_latent_len]) 99 | print("load from 2") 100 | 101 | assert measurements.shape[0] == finetune_latent.shape[0],"some data are corrupted" 102 | sample_num = finetune_latent.shape[0] 103 | texel_sequence = np.arange(sample_num) 104 | 105 | pf_result = open(data_root+f"pass3_latent_{args.tex_resolution}.bin","wb") 106 | pf_result_lumi = open(data_root+"finetune_lumi.bin","wb") 107 | 108 | optimize_step = 300 if args.if_continue else 500 109 | 110 | lr = 0.02 111 | ptr = 0 112 | while True: 113 | if ptr % 30000 == 0: 114 | start = datetime.now() 115 | print(f"PASS 3 [{which_test_thread}]/{ptr}/{sample_num} {start}") 116 | 117 | tmp_sequence = texel_sequence[ptr:ptr+args.batchsize] 118 | tmp_seq_size = tmp_sequence.shape[0] 119 | 120 | if tmp_seq_size == 0: 121 | break 122 | 123 | tmp_measurements = measurements[tmp_sequence] 124 | 125 | tmp_x_guess = torch.cuda.FloatTensor(finetune_latent[tmp_sequence],device=compute_device) 126 | tmp_x_guess = Variable(tmp_x_guess,requires_grad=True) 127 | optimizer = torch.optim.Adam([tmp_x_guess,], lr = lr) 128 | 129 | loss_step = [] 130 | loss_precent = [] 131 | for opt_step in range(optimize_step): 132 | 133 | color_latent = tmp_x_guess[:,:3 * args.color_latent_len].reshape([tmp_seq_size,3,args.color_latent_len]) 134 | shape_latent = tmp_x_guess[:,3 * args.color_latent_len:].reshape([tmp_seq_size,1,args.shape_latent_len]).repeat(1,3,1) 135 | 136 | color_shape_latent = torch.cat([color_latent,shape_latent],dim=-1) 137 | color_shape_latent = color_shape_latent.reshape([tmp_seq_size*3,latent_len]) 138 | 139 | _,tmp_nn_lumi = inference_net(color_shape_latent,input_is_latent=True) 140 | tmp_nn_lumi = torch.max(torch.zeros_like(tmp_nn_lumi),tmp_nn_lumi) 141 | tmp_nn_lumi = tmp_nn_lumi.reshape([tmp_seq_size,3,-1]).permute(0,2,1).unsqueeze(dim=1) 142 | 143 | tmp_lumi_measurements = torch.sum(lighting_patterns*tmp_nn_lumi,dim=2).reshape([tmp_seq_size,args.finetune_use_num,3]) 144 | 145 | loss = torch.nn.functional.mse_loss(torch.pow(tmp_measurements,1/2), torch.pow(tmp_lumi_measurements,1/2),reduction='sum') 146 | 147 | optimizer.zero_grad() 148 | loss.backward() 149 | optimizer.step() 150 | 151 | loss_step.append(loss) 152 | if opt_step >= 10 and not args.if_continue: 153 | tmp_loss_percent = loss/loss_step[opt_step-10]*100 154 | loss_precent.append(tmp_loss_percent) 155 | if len(loss_precent) >= 10: 156 | if torch.mean(torch.stack(loss_precent[-5:],dim=0)) > 98.0 and opt_step > 75: 157 | break 158 | 159 | 160 | tmp_x_guess.detach().cpu().numpy().astype(np.float32).tofile(pf_result) 161 | 162 | ptr += args.batchsize 163 | 164 | 165 | pf_result.close() 166 | pf_result_lumi.close() 167 | 168 | print(which_test_thread, "done.") 169 | -------------------------------------------------------------------------------- /data_processing/finetune/fitting_master.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the experimental code for paper ``Ma et al., OpenSVBRDF: A Database of Measured Spatially-Varying Reflectance, SIGGRAPH Asia 2023``. 3 | This script is suboptimal and experimental. 4 | There may be redundant lines and functionalities. 5 | 6 | Xiaohe Ma, 2024/02 7 | ''' 8 | 9 | ''' 10 | This script runs the three steps of the per-pixel fine-tuning process. 11 | 12 | ''' 13 | 14 | import os 15 | from subprocess import Popen 16 | import argparse 17 | import queue 18 | from warnings import simplefilter 19 | simplefilter(action='ignore', category=FutureWarning) 20 | import warnings 21 | warnings.filterwarnings('ignore') 22 | 23 | exclude_gpu_list = [] 24 | 25 | if __name__ == "__main__": 26 | parser = argparse.ArgumentParser(usage="fitting using multi cards") 27 | 28 | parser.add_argument("data_root") 29 | parser.add_argument("thread_num",type=int) 30 | parser.add_argument("server_num",type=int) 31 | parser.add_argument("which_server",type=int) 32 | parser.add_argument("if_dump",type=int) 33 | parser.add_argument("lighting_pattern_num",type=int) 34 | parser.add_argument("finetune_use_num",type=int) 35 | parser.add_argument("tex_resolution",type=int) 36 | parser.add_argument("main_cam_id",type=int) 37 | parser.add_argument("model_file",type=str) 38 | parser.add_argument("pattern_file",type=str) 39 | parser.add_argument("shape_latent_len",type=int) 40 | parser.add_argument("color_latent_len",type=int) 41 | 42 | parser.add_argument("save_lumi",type=int) 43 | parser.add_argument("step",type=str) 44 | parser.add_argument("--if_continue",type=int,default=0) 45 | 46 | args = parser.parse_args() 47 | 48 | deviceIDs = [0,1,2,3] 49 | gpu_num = len(deviceIDs) 50 | print("available gpu num:",gpu_num) 51 | pool = [] 52 | p_log_f = open(args.data_root+"fitting_log_all.txt","w",buffering=1) 53 | ################################## 54 | # step 1 55 | ################################## 56 | if '1' in args.step: 57 | thread_per_gpu = [0]*len(deviceIDs) 58 | q = queue.Queue() 59 | [q.put(i) for i in range(0, args.server_num*args.thread_num)] 60 | try_num = [0]*(args.thread_num*args.server_num) 61 | pool = [] 62 | while not q.empty(): 63 | which_thread = q.get() 64 | try_num[which_thread] = try_num[which_thread]+1 65 | print("starting thread:{}".format(which_thread)) 66 | which_gpu = which_thread%gpu_num 67 | my_env = os.environ.copy() 68 | 69 | this_args = [ 70 | "python", 71 | "finetune_pass1.py", 72 | "--data_for_server_root", 73 | "{}".format(args.data_root), 74 | "--lighting_pattern_num", 75 | "{}".format(args.lighting_pattern_num), 76 | "--finetune_use_num", 77 | "{}".format(args.finetune_use_num), 78 | "--thread_ids", 79 | "{}".format(which_thread), 80 | "--gpu_id", 81 | "{}".format(deviceIDs[which_gpu]), 82 | "--tex_resolution", 83 | "{}".format(args.tex_resolution), 84 | "--total_thread_num", 85 | "{}".format(args.thread_num*args.server_num), 86 | "--main_cam_id", 87 | "{}".format(args.main_cam_id), 88 | "--model_file", 89 | "{}".format(args.model_file), 90 | "--pattern_file", 91 | "{}".format(args.pattern_file), 92 | "--shape_latent_len", 93 | "{}".format(args.shape_latent_len), 94 | "--color_latent_len", 95 | "{}".format(args.color_latent_len) 96 | ] 97 | if args.if_dump: 98 | this_args.append("--need_dump") 99 | if args.save_lumi: 100 | this_args.append("--save_lumi") 101 | theProcess = Popen( 102 | this_args, 103 | env=my_env 104 | ) 105 | pool.append(theProcess) 106 | thread_per_gpu[which_gpu] += 1 107 | if thread_per_gpu.count(3) == len(thread_per_gpu): 108 | exit_codes = [p.wait() for p in pool] 109 | print("fitting rhod rhos exit codes:",exit_codes) 110 | for i in range(len(exit_codes)): 111 | crash_thread_id = int(pool[i].args[9]) 112 | if exit_codes[i] != 0 and try_num[crash_thread_id] < 2: 113 | print("thread:{} crashed!".format(crash_thread_id)) 114 | q.put(crash_thread_id) 115 | pool = [] 116 | thread_per_gpu = [0]*len(deviceIDs) 117 | exit_codes = [ p.wait() for p in pool ] 118 | print("fitting rhod rhos exit codes:",exit_codes) 119 | p_log_f.write("fitting rhod rhos exit codes:{}\n".format(exit_codes)) 120 | 121 | ################################## 122 | #step 2 123 | ################################## 124 | if '2' in args.step: 125 | thread_per_gpu = [0]*len(deviceIDs) 126 | q = queue.Queue() 127 | [q.put(i) for i in range(0, args.server_num*args.thread_num)] 128 | 129 | try_num = [0]*(args.thread_num*args.server_num) 130 | pool = [] 131 | while not q.empty(): 132 | which_thread = q.get() 133 | try_num[which_thread] = try_num[which_thread]+1 134 | print("starting thread:{}".format(which_thread)) 135 | which_gpu = which_thread%gpu_num 136 | my_env = os.environ.copy() 137 | 138 | this_args = [ 139 | "python", 140 | "finetune_pass2.py", 141 | "--data_for_server_root", 142 | "{}".format(args.data_root), 143 | "--lighting_pattern_num", 144 | "{}".format(args.lighting_pattern_num), 145 | "--finetune_use_num", 146 | "{}".format(args.finetune_use_num), 147 | "--thread_ids", 148 | "{}".format(which_thread), 149 | "--gpu_id", 150 | "{}".format(deviceIDs[which_gpu]), 151 | "--tex_resolution", 152 | "{}".format(args.tex_resolution), 153 | "--total_thread_num", 154 | "{}".format(args.thread_num*args.server_num), 155 | "--main_cam_id", 156 | "{}".format(args.main_cam_id), 157 | "--model_file", 158 | "{}".format(args.model_file), 159 | "--pattern_file", 160 | "{}".format(args.pattern_file), 161 | "--shape_latent_len", 162 | "{}".format(args.shape_latent_len), 163 | "--color_latent_len", 164 | "{}".format(args.color_latent_len) 165 | ] 166 | if args.if_dump: 167 | this_args.append("--need_dump") 168 | if args.save_lumi: 169 | this_args.append("--save_lumi") 170 | theProcess = Popen( 171 | this_args, 172 | env=my_env 173 | ) 174 | 175 | pool.append(theProcess) 176 | thread_per_gpu[which_gpu] += 1 177 | if thread_per_gpu.count(3) == len(thread_per_gpu): 178 | exit_codes = [p.wait() for p in pool] 179 | print("fitting rhod rhos exit codes:",exit_codes) 180 | for i in range(len(exit_codes)): 181 | crash_thread_id = int(pool[i].args[9]) 182 | if exit_codes[i] != 0 and try_num[crash_thread_id] < 2: 183 | print("thread:{} crashed!".format(crash_thread_id)) 184 | q.put(crash_thread_id) 185 | pool = [] 186 | thread_per_gpu = [0]*len(deviceIDs) 187 | exit_codes = [ p.wait() for p in pool ] 188 | print("fitting rhod rhos exit codes:",exit_codes) 189 | p_log_f.write("fitting rhod rhos exit codes:{}\n".format(exit_codes)) 190 | 191 | ################################## 192 | #step 3 193 | ################################## 194 | if '3' in args.step: 195 | thread_per_gpu = [0]*len(deviceIDs) 196 | q = queue.Queue() 197 | [q.put(i) for i in range(0, args.server_num*args.thread_num)] 198 | 199 | try_num = [0]*(args.thread_num*args.server_num) 200 | pool = [] 201 | while not q.empty(): 202 | which_thread = q.get() 203 | try_num[which_thread] = try_num[which_thread]+1 204 | print("starting thread:{}".format(which_thread)) 205 | which_gpu = which_thread%gpu_num 206 | my_env = os.environ.copy() 207 | 208 | this_args = [ 209 | "python", 210 | "finetune_pass3.py", 211 | "--data_for_server_root", 212 | "{}".format(args.data_root), 213 | "--lighting_pattern_num", 214 | "{}".format(args.lighting_pattern_num), 215 | "--finetune_use_num", 216 | "{}".format(args.finetune_use_num), 217 | "--thread_ids", 218 | "{}".format(which_thread), 219 | "--gpu_id", 220 | "{}".format(deviceIDs[which_gpu]), 221 | "--tex_resolution", 222 | "{}".format(args.tex_resolution), 223 | "--total_thread_num", 224 | "{}".format(args.thread_num*args.server_num), 225 | "--main_cam_id", 226 | "{}".format(args.main_cam_id), 227 | "--model_file", 228 | "{}".format(args.model_file), 229 | "--pattern_file", 230 | "{}".format(args.pattern_file), 231 | "--shape_latent_len", 232 | "{}".format(args.shape_latent_len), 233 | "--color_latent_len", 234 | "{}".format(args.color_latent_len), 235 | "--if_continue", 236 | "{}".format(args.if_continue) 237 | ] 238 | if args.if_dump: 239 | this_args.append("--need_dump") 240 | if args.save_lumi: 241 | this_args.append("--save_lumi") 242 | theProcess = Popen( 243 | this_args, 244 | env=my_env 245 | ) 246 | 247 | pool.append(theProcess) 248 | thread_per_gpu[which_gpu] += 1 249 | if thread_per_gpu.count(2) == len(thread_per_gpu): 250 | exit_codes = [p.wait() for p in pool] 251 | print("fitting rhod rhos exit codes:",exit_codes) 252 | for i in range(len(exit_codes)): 253 | crash_thread_id = int(pool[i].args[9]) 254 | if exit_codes[i] != 0 and try_num[crash_thread_id] < 2: 255 | print("thread:{} crashed!".format(crash_thread_id)) 256 | q.put(crash_thread_id) 257 | pool = [] 258 | thread_per_gpu = [0]*len(deviceIDs) 259 | exit_codes = [ p.wait() for p in pool ] 260 | print("fitting rhod rhos exit codes:",exit_codes) 261 | p_log_f.write("fitting rhod rhos exit codes:{}\n".format(exit_codes)) 262 | -------------------------------------------------------------------------------- /data_processing/finetune/gather_finetune_results.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the experimental code for paper ``Ma et al., OpenSVBRDF: A Database of Measured Spatially-Varying Reflectance, SIGGRAPH Asia 2023``. 3 | This script is suboptimal and experimental. 4 | There may be redundant lines and functionalities. 5 | 6 | Xiaohe Ma, 2024/02 7 | ''' 8 | 9 | 10 | ''' 11 | This script gathers finetuning results and organizes documents for further operation. 12 | 13 | by Leyao 14 | ''' 15 | import numpy as np 16 | import cv2 17 | import os 18 | import OpenEXR, array 19 | import Imath 20 | from sklearn.decomposition import PCA 21 | import argparse 22 | 23 | FLOAT = Imath.Channel(Imath.PixelType(Imath.PixelType.FLOAT)) 24 | 25 | if __name__ == "__main__": 26 | parser = argparse.ArgumentParser(usage="fitting using multi card") 27 | 28 | parser.add_argument("data_root") 29 | parser.add_argument("save_root") 30 | parser.add_argument("thread_num",type=int) 31 | parser.add_argument("server_num",type=int) 32 | parser.add_argument("tex_resolution",type=int) 33 | parser.add_argument("--latent_len",type=int,default=72) 34 | 35 | args = parser.parse_args() 36 | os.makedirs(args.save_root, exist_ok=True) 37 | 38 | pf_latent = {} 39 | for latent_id in range(4): 40 | pf_latent[f"{latent_id}"] = open(args.save_root+f"pass{latent_id}_latent_{args.tex_resolution}.bin","wb") 41 | 42 | for latent_id in range(4): 43 | for which_thread in range(args.thread_num*args.server_num): 44 | tmp_latent = np.fromfile(args.data_root+f"{which_thread}/pass{latent_id}_latent_{args.tex_resolution}.bin", np.float32) 45 | tmp_latent.astype(np.float32).tofile(pf_latent[f"{latent_id}"]) 46 | 47 | for latent_id in range(4): 48 | pf_latent[f"{latent_id}"].close() 49 | 50 | data = np.fromfile(args.save_root+f"pass3_latent_{args.tex_resolution}.bin",np.float32).reshape([args.tex_resolution,args.tex_resolution,-1]) 51 | 52 | print(data.shape) 53 | 54 | header = OpenEXR.Header(args.tex_resolution,args.tex_resolution) 55 | 56 | channels_num = data.shape[2] 57 | new_channels = {} 58 | 59 | for which_channel in range(data.shape[2]): 60 | tmpdata = data[:,:,which_channel].reshape([args.tex_resolution,args.tex_resolution,1]).tobytes() 61 | new_channels['L_{}'.format(which_channel)] = tmpdata 62 | 63 | header['channels'] = dict([(c, FLOAT) for c in new_channels.keys()]) 64 | header['compression'] = Imath.Compression(Imath.Compression.PIZ_COMPRESSION) 65 | out = OpenEXR.OutputFile(args.save_root+"latent.exr", header) 66 | out.writePixels(new_channels) 67 | 68 | data = data.reshape([-1, args.latent_len]) 69 | pca = PCA(n_components=3) 70 | pca.fit(data) 71 | data_reduction = pca.transform(data) 72 | data_reduction_min = np.min(data_reduction) 73 | data_reduction_max = np.max(data_reduction) 74 | data_reduction_delta = data_reduction_max - data_reduction_min 75 | data_reduction = (data_reduction - data_reduction_min) / data_reduction_delta 76 | data_reduction = data_reduction.reshape([args.tex_resolution,args.tex_resolution,3]) 77 | 78 | cv2.imwrite(args.save_root+"latent_pca.png",data_reduction*255) 79 | 80 | -------------------------------------------------------------------------------- /data_processing/finetune/infer.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the experimental code for paper ``Ma et al., OpenSVBRDF: A Database of Measured Spatially-Varying Reflectance, SIGGRAPH Asia 2023``. 3 | This script is suboptimal and experimental. 4 | There may be redundant lines and functionalities. 5 | 6 | Xiaohe Ma, 2024/02 7 | ''' 8 | 9 | ''' 10 | This script infers input data by loading PlanarScannerNet using a pre-trained neural network model. 11 | ''' 12 | 13 | import torch 14 | import argparse 15 | import random 16 | import sys 17 | import numpy as np 18 | import os 19 | 20 | import planar_scanner_net_inference 21 | 22 | if __name__ == "__main__": 23 | 24 | parser = argparse.ArgumentParser() 25 | parser.add_argument("data_root") 26 | parser.add_argument("model_root") 27 | parser.add_argument("model_file_name") 28 | parser.add_argument("lighting_pattern_num",type=int) 29 | parser.add_argument("m_len",type=int) 30 | parser.add_argument("tex_resolution",type=int) 31 | parser.add_argument("--training_gpu",type=int,default=0) 32 | parser.add_argument("--batch_size",type=int,default=100) 33 | parser.add_argument("--cam_num",type=int,default=2) 34 | parser.add_argument("--main_cam_id",type=int,default=0) 35 | 36 | parser.add_argument("--layers",type=int,default=8) 37 | parser.add_argument("--shape_latent_len",type=int,default=2) 38 | parser.add_argument("--color_latent_len",type=int,default=2) 39 | parser.add_argument("--need_dump",action="store_true") 40 | 41 | 42 | args = parser.parse_args() 43 | 44 | latent_len = args.shape_latent_len + args.color_latent_len 45 | 46 | # Define configuration parameters 47 | train_configs = {} 48 | train_configs["rendering_devices"] = [torch.device("cuda:{}".format(args.training_gpu))] # for multiple GPU 49 | train_configs["training_device"] = torch.device("cuda:{}".format(args.training_gpu)) 50 | train_configs["layers"] = args.layers 51 | train_configs["lighting_pattern_num"] = args.lighting_pattern_num 52 | train_configs["m_len"] = args.m_len 53 | train_configs["train_lighting_pattern"] = False 54 | train_configs["lumitexel_length"] = 64*64*3 55 | train_configs["cam_num"] = args.cam_num 56 | train_configs["main_cam_id"] = args.main_cam_id 57 | train_configs["latent_len"] = latent_len 58 | train_configs["color_latent_len"] = args.color_latent_len 59 | train_configs["shape_latent_len"] = args.shape_latent_len 60 | train_configs["data_root"] = args.data_root 61 | train_configs["batch_size"] = args.batch_size*3 62 | train_configs["pre_load_buffer_size"] = 500000 63 | 64 | save_root = args.data_root + f"texture_{args.tex_resolution}/" 65 | 66 | os.makedirs(save_root,exist_ok=True) 67 | 68 | # Load the pre-trained model 69 | model = planar_scanner_net_inference.PlanarScannerNet(train_configs) 70 | inference_device = train_configs["training_device"] 71 | pretrained_dict = torch.load(args.model_root + args.model_file_name, map_location='cuda:0') 72 | 73 | print("loading trained model...") 74 | something_not_found = False 75 | model_dict = model.state_dict() 76 | for k,_ in model_dict.items(): 77 | if k not in pretrained_dict: 78 | print("not found:", k) 79 | something_not_found = True 80 | if something_not_found: 81 | exit() 82 | model_dict = model.state_dict() 83 | pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} 84 | model_dict.update(pretrained_dict) 85 | model.load_state_dict(model_dict) 86 | model.to(inference_device) 87 | 88 | 89 | model.eval() 90 | pf_map = {} 91 | pf_map["latent"] = open(save_root+f"latent_{args.tex_resolution}.bin","wb") 92 | 93 | record_size_byte = args.cam_num * args.m_len * args.lighting_pattern_num * 2 * 4 94 | 95 | if args.m_len == 1: 96 | record_size_byte *= 3 97 | pf_measurements = open(args.data_root+f"texture_{args.tex_resolution}/measurements_{args.tex_resolution}.bin", "rb") 98 | pf_measurements.seek(0,2) 99 | texel_num = pf_measurements.tell()//record_size_byte 100 | pf_measurements.seek(0,0) 101 | print("texel num : ", texel_num) 102 | texel_sequence = np.arange(texel_num) 103 | start_ptr = 0 104 | ptr = start_ptr 105 | batch_size = args.batch_size 106 | lumitexel_length = train_configs["lumitexel_length"] 107 | lighting_pattern_num = args.lighting_pattern_num 108 | 109 | # Process data in batches 110 | while True: 111 | tmp_sequence = texel_sequence[ptr:ptr+batch_size] 112 | if tmp_sequence.shape[0] == 0: 113 | break 114 | tmp_seq_size = tmp_sequence.shape[0] 115 | 116 | # Read raw measurements 117 | tmp_measurements_raw = np.fromfile(pf_measurements,np.float32,count=record_size_byte//4*tmp_seq_size).reshape([tmp_seq_size,args.cam_num,lighting_pattern_num*2,3]) 118 | tmp_measurements_raw = torch.from_numpy(tmp_measurements_raw).to(inference_device) 119 | 120 | tmp_measurements_raw_mean = torch.mean(tmp_measurements_raw,dim=-1) 121 | 122 | tmp_measurements_raw_mean = tmp_measurements_raw_mean[:,:,::2] - tmp_measurements_raw_mean[:,:,1::2] 123 | 124 | # Perform inference 125 | res = model(tmp_measurements_raw_mean) 126 | latent = res["nn_latent"] 127 | 128 | latent.astype(np.float32).tofile(pf_map["latent"]) 129 | 130 | 131 | ptr += batch_size 132 | 133 | for a_key in pf_map: 134 | pf_map[a_key].close() 135 | -------------------------------------------------------------------------------- /data_processing/finetune/latent_net.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import numpy as np 4 | from collections import OrderedDict 5 | 6 | class LatentNet(nn.Module): 7 | ''' 8 | LatentNet defines a neural network model for extracting latent features of color and shape. 9 | 10 | Functions: 11 | color_latent_decoder: construct a neural network model for latent features of color. 12 | 13 | shape_latent_decoder: construct a neural network model for latent features of shape. 14 | 15 | ''' 16 | def __init__(self,args): 17 | super(LatentNet,self).__init__() 18 | self.leaf_nodes_num = args["leaf_nodes_num"] 19 | self.cam_num = args["cam_num"] 20 | self.input_length = args["lighting_pattern_num"] * self.cam_num * self.leaf_nodes_num 21 | 22 | 23 | self.color_latent_len = args["color_latent_len"] 24 | self.shape_latent_len = args["shape_latent_len"] 25 | self.latent_len = self.shape_latent_len + self.color_latent_len 26 | 27 | # construct model 28 | self.shape_layer_width = [256,1024,3072,1024,256,self.shape_latent_len] 29 | self.shape_layer_width = [value * self.leaf_nodes_num for value in self.shape_layer_width] 30 | 31 | self.color_layer_width = [256,1024,512,64,self.color_latent_len] 32 | self.color_layer_width = [value * self.leaf_nodes_num for value in self.color_layer_width] 33 | 34 | self.shape_latent_net_model = self.shape_latent_decoder(self.input_length) 35 | self.color_latent_net_model = self.color_latent_decoder(self.input_length) 36 | 37 | def color_latent_decoder(self,input_size,name_prefix = "color_latent_net_"): 38 | layer_stack = OrderedDict() 39 | 40 | layer_count = 0 41 | input_size = self.input_length 42 | 43 | for which_layer in self.color_layer_width[:-1]: 44 | output_size = which_layer 45 | layer_stack[name_prefix+"Conv_{}".format(layer_count)] = torch.nn.Conv1d(input_size, output_size, kernel_size=1,groups=self.leaf_nodes_num) 46 | if which_layer > 1: 47 | layer_stack[name_prefix+"BN_{}".format(layer_count)] = nn.BatchNorm1d(output_size) 48 | layer_stack[name_prefix+"LeakyRelu_{}".format(layer_count)] = nn.LeakyReLU() 49 | layer_count+=1 50 | input_size = output_size 51 | 52 | output_size = self.color_layer_width[-1] 53 | layer_stack["Linear_{}".format(layer_count)] = torch.nn.Conv1d(input_size, output_size, kernel_size=1, groups=self.leaf_nodes_num) 54 | 55 | layer_stack = nn.Sequential(layer_stack) 56 | 57 | return layer_stack 58 | 59 | def shape_latent_decoder(self,input_size,name_prefix = "shape_latent_net_"): 60 | layer_stack = OrderedDict() 61 | 62 | layer_count = 0 63 | input_size = self.input_length 64 | 65 | for which_layer in self.shape_layer_width[:-1]: 66 | output_size = which_layer 67 | layer_stack[name_prefix+"Conv_{}".format(layer_count)] = torch.nn.Conv1d(input_size, output_size, kernel_size=1,groups=self.leaf_nodes_num) 68 | if which_layer > 2: 69 | layer_stack[name_prefix+"BN_{}".format(layer_count)] = nn.BatchNorm1d(output_size) 70 | layer_stack[name_prefix+"LeakyRelu_{}".format(layer_count)] = nn.LeakyReLU() 71 | layer_count+=1 72 | input_size = output_size 73 | 74 | output_size = self.shape_layer_width[-1] 75 | layer_stack["Linear_{}".format(layer_count)] = torch.nn.Conv1d(input_size, output_size, kernel_size=1, groups=self.leaf_nodes_num) 76 | 77 | layer_stack = nn.Sequential(layer_stack) 78 | 79 | return layer_stack 80 | 81 | 82 | def forward(self,measurements): 83 | batch_size = measurements.size()[0] 84 | x_n = measurements.unsqueeze(dim=1).unsqueeze(dim=-1).repeat(1,self.leaf_nodes_num,1,1).reshape([batch_size,-1,1]) 85 | 86 | 87 | shape_latent = self.shape_latent_net_model(x_n) 88 | color_latent = self.color_latent_net_model(x_n) 89 | 90 | 91 | latent = torch.cat([color_latent,shape_latent],dim=1) 92 | latent = latent.reshape([batch_size,self.leaf_nodes_num,self.latent_len]) 93 | 94 | 95 | return latent -------------------------------------------------------------------------------- /data_processing/finetune/planar_scanner_net_inference.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch.utils.data 3 | import numpy as np 4 | import math 5 | 6 | from choose_branch_net import ChooseBranchNet 7 | 8 | from latent_net import LatentNet 9 | 10 | class PlanarScannerNet(nn.Module): 11 | ''' 12 | PlanarScannerNet implements a neural network model for finetuning. 13 | 14 | Functions: 15 | precompute_table: define a table for searching on initialization. 16 | 17 | forward: forward propagation method for input data. 18 | 19 | ''' 20 | def __init__(self, args): 21 | super(PlanarScannerNet, self).__init__() 22 | ######################################## 23 | ## parse configuration ### 24 | ######################################## 25 | self.args = args 26 | self.training_device = args["training_device"] 27 | self.layers = args["layers"] 28 | 29 | self.m_len = args["m_len"] 30 | self.cam_num = args["cam_num"] 31 | self.main_cam_id = args["main_cam_id"] 32 | self.batch_size = args["batch_size"] 33 | self.sublings_num = 2 34 | self.lumitexel_length = args["lumitexel_length"] 35 | ######################################## 36 | ## loading setup configuration ### 37 | ######################################## 38 | self.leaf_nodes_num = int(math.pow(2,self.layers-1)) 39 | args["classifier_num"] = self.leaf_nodes_num - 1 40 | args["leaf_nodes_num"] = self.leaf_nodes_num 41 | 42 | self.ptr = self.leaf_nodes_num - 1 43 | self.return_all_leaf = False 44 | ######################################## 45 | ## define net modules ### 46 | ######################################## 47 | self.l2_loss_fn = torch.nn.MSELoss(reduction='sum') 48 | 49 | self.latent_net = LatentNet(args) 50 | self.choose_branch_net = ChooseBranchNet(args) 51 | 52 | self.search_table = self.precompute_table() 53 | 54 | self.leaf_indices = torch.arange(self.leaf_nodes_num).long() 55 | self.layer_indices = self.search_table[:,:,0].long() 56 | self.base = torch.from_numpy(np.array([16,8,4,2,1])).long().to(self.training_device) 57 | 58 | def precompute_table(self): 59 | search_table = [] 60 | for which_node in range(self.ptr, self.ptr+self.leaf_nodes_num): 61 | tmp_router = [] 62 | subling = which_node 63 | parent = (which_node - 1) // self.sublings_num 64 | while parent >= 0: 65 | tmp_router.append(parent) 66 | tmp_router.append((subling-1)%self.sublings_num) 67 | subling = parent 68 | parent = (parent-1) // self.sublings_num 69 | 70 | search_table.append(torch.as_tensor(tmp_router)) 71 | 72 | search_table = torch.stack(search_table,dim=0).reshape([self.leaf_nodes_num,self.layers-1,2]) 73 | 74 | return search_table 75 | 76 | def calculate_loss(self,pred,label): 77 | tmp_loss = pred - label 78 | tmp_loss = torch.sum(tmp_loss*tmp_loss,dim=1,keepdims=True).reshape(self.batch_size,1) 79 | return tmp_loss.clone() 80 | 81 | 82 | def forward(self, batch_data,call_type="train"): 83 | measurements = batch_data 84 | batch_size = measurements.shape[0] 85 | self.batch_indices = torch.arange(batch_size)[:,None] 86 | 87 | measurements_main_cam = measurements[:,self.main_cam_id] 88 | 89 | expand_measurements = measurements_main_cam.unsqueeze(dim=1).repeat(1,self.layers-1,1).reshape([batch_size,-1]).unsqueeze(dim=-1) 90 | measurements = torch.cat([measurements[:,0],measurements[:,1]],dim=-1) 91 | 92 | clean_logits = self.choose_branch_net(expand_measurements) 93 | 94 | expand_weight = [] 95 | for which_layer in range(self.layers-1): 96 | tmp_choose_net_output = clean_logits[:,which_layer].unsqueeze(dim=1).repeat(1,int(math.pow(self.sublings_num,which_layer)),1) 97 | 98 | expand_weight.append(tmp_choose_net_output) 99 | 100 | all_classifier_weight = torch.cat(expand_weight,dim=1) 101 | 102 | 103 | all_leaf_index = self.search_table.clone() 104 | all_leaf_index = all_leaf_index.repeat(batch_size,1,1,1) 105 | 106 | all_classifier_weight = all_classifier_weight.reshape([batch_size,-1,2]).unsqueeze(dim=1).repeat(1,self.leaf_nodes_num,1,1) 107 | 108 | all_leaf_weight = all_classifier_weight.reshape([batch_size*self.leaf_nodes_num,-1,2]) 109 | 110 | all_leaf_index = all_leaf_index.reshape([batch_size*self.leaf_nodes_num,-1,2]) 111 | 112 | self.batch_indices = torch.arange(batch_size*self.leaf_nodes_num)[:,None] 113 | 114 | all_leaf_weight = all_leaf_weight[self.batch_indices,all_leaf_index[:,:,0]] 115 | 116 | 117 | all_leaf_weight = all_leaf_weight.reshape([-1,2]) 118 | 119 | all_leaf_index = all_leaf_index[:,:,1].reshape([-1,1]) 120 | self.batch_indices = torch.arange(all_leaf_index.shape[0])[:,None] 121 | 122 | all_leaf_weight = all_leaf_weight[self.batch_indices,all_leaf_index].reshape([batch_size,self.leaf_nodes_num,self.layers-1]) 123 | 124 | all_leaf_weight = torch.prod(all_leaf_weight,dim=2) 125 | 126 | 127 | gates = all_leaf_weight 128 | 129 | ################### BALANCE LOSS #################### 130 | 131 | router = torch.argsort(gates,dim=-1)[:,-1] 132 | 133 | router = router.unsqueeze(dim=-1) 134 | 135 | nn_latent = self.latent_net(measurements) 136 | nn_latent = nn_latent.reshape([batch_size,self.leaf_nodes_num,-1]) 137 | 138 | select_nn_latent = torch.zeros_like(nn_latent[:,0]) 139 | 140 | for which_node in range(self.leaf_nodes_num): 141 | 142 | tmp_router_idx = torch.where(router == which_node)[0] 143 | if tmp_router_idx.shape[0] > 0: 144 | select_nn_latent[tmp_router_idx] = nn_latent[tmp_router_idx][:,which_node] 145 | 146 | 147 | term_map = { 148 | "nn_latent":select_nn_latent.detach().cpu().numpy(), 149 | } 150 | 151 | return term_map 152 | -------------------------------------------------------------------------------- /data_processing/finetune/run.sh: -------------------------------------------------------------------------------- 1 | for variable in 18 2 | do 3 | class_name=paper 4 | formatted_variable=$(printf "%04d" $variable) 5 | name=$class_name$formatted_variable 6 | tex_resolution=1024 7 | 8 | data_root=../../database_data/$name/output/ 9 | model_root=../../database_model/ 10 | model_file_name=model_state_450000.pkl 11 | 12 | lighting_pattern_num=32 13 | measurement_len=1 14 | shape_latent_len=48 15 | color_latent_len=8 16 | python infer.py $data_root $model_root $model_file_name $lighting_pattern_num $measurement_len $tex_resolution --layers 6 --shape_latent_len $shape_latent_len --color_latent_len $color_latent_len 17 | 18 | if_dump=1 19 | save_lumi=1 20 | thread_num=4 21 | server_num=4 22 | which_server=0 23 | lighting_pattern_num=64 24 | main_cam_id=0 25 | model_file_name=$model_root/latent_48_24_500000_2437925.pkl 26 | 27 | pattern_file_name=$model_root/opt_W_64.bin 28 | 29 | finetune_use_num=64 30 | step="123" 31 | if_continue=0 32 | 33 | python split_data_and_prepare_for_server.py $data_root $lighting_pattern_num $thread_num $server_num $which_server $tex_resolution $main_cam_id $shape_latent_len $color_latent_len 34 | 35 | save_root=$data_root/latent/ 36 | data_root=$data_root/data_for_server/ 37 | 38 | python fitting_master.py $data_root $thread_num $server_num $which_server $if_dump $lighting_pattern_num $finetune_use_num $tex_resolution $main_cam_id $model_file_name $pattern_file_name $shape_latent_len $color_latent_len $save_lumi $step --if_continue $if_continue 39 | 40 | python gather_finetune_results.py $data_root $save_root $thread_num $server_num $tex_resolution 41 | 42 | done -------------------------------------------------------------------------------- /data_processing/finetune/split_data_and_prepare_for_server.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the experimental code for paper ``Ma et al., OpenSVBRDF: A Database of Measured Spatially-Varying Reflectance, SIGGRAPH Asia 2023``. 3 | This script is suboptimal and experimental. 4 | There may be redundant lines and functionalities. 5 | 6 | Xiaohe Ma, 2024/02 7 | ''' 8 | 9 | 10 | ''' 11 | This script splits the data and prepares for the server. 12 | 13 | ''' 14 | 15 | import numpy as np 16 | import os 17 | import argparse 18 | 19 | parser = argparse.ArgumentParser(usage="split relighting brdf not slice") 20 | 21 | parser.add_argument("data_root") 22 | parser.add_argument("lighting_pattern_num",type=int) 23 | parser.add_argument("thread_num",type=int) 24 | parser.add_argument("server_num",type=int) 25 | parser.add_argument("which_server",type=int) 26 | parser.add_argument("tex_resolution",type=int) 27 | parser.add_argument("main_cam_id",type=int) 28 | parser.add_argument("shape_latent_len",type=int) 29 | parser.add_argument("color_latent_len",type=int) 30 | 31 | args = parser.parse_args() 32 | 33 | if __name__ == "__main__": 34 | data_root = args.data_root + f"texture_{args.tex_resolution}/" 35 | target_root = args.data_root + "data_for_server/" 36 | 37 | latent_len = args.shape_latent_len + args.color_latent_len 38 | 39 | m_len_perview = 3 40 | total_thread_num = args.thread_num*args.server_num 41 | 42 | os.makedirs(target_root, exist_ok=True) 43 | 44 | pf_latent = open(data_root+f"latent_{args.tex_resolution}.bin") 45 | pf_latent.seek(0,2) 46 | pixel_num = pf_latent.tell() //4//latent_len 47 | print("[SPLITTER]pixel num:",pixel_num) 48 | pf_latent.seek(0,0) 49 | 50 | pf_pos = open(data_root+"positions.bin") 51 | 52 | pf_measurement = open(data_root+f"line_measurements_{args.tex_resolution}.bin") 53 | 54 | num_per_thread = int(pixel_num//total_thread_num) 55 | 56 | ptr = 0 57 | for thread_id in range(total_thread_num): 58 | tmp_dir = target_root+f"{thread_id}/" 59 | os.makedirs(tmp_dir, exist_ok=True) 60 | 61 | cur_batchsize = num_per_thread if (not thread_id == total_thread_num-1) else (pixel_num-(total_thread_num-1)*num_per_thread) 62 | 63 | tmp_latents = np.fromfile(pf_latent,np.float32,cur_batchsize*latent_len) 64 | tmp_latents.astype(np.float32).tofile(tmp_dir+f"latent_{args.tex_resolution}.bin") 65 | tmp_positions = np.fromfile(pf_pos,np.float32,cur_batchsize*3) 66 | tmp_positions.astype(np.float32).tofile(tmp_dir+"positions.bin") 67 | 68 | 69 | tmp_measurements = np.fromfile(pf_measurement,np.float32, cur_batchsize*args.lighting_pattern_num*m_len_perview) 70 | tmp_measurements.astype(np.float32).tofile(tmp_dir+f"gt_measurements_{args.tex_resolution}.bin") 71 | print("thread:",thread_id," num:",cur_batchsize) 72 | 73 | ptr += cur_batchsize 74 | 75 | remain_data = np.fromfile(pf_measurement,np.uint8) 76 | 77 | if len(remain_data) > 0: 78 | print("meaurements file is not at the end!") 79 | exit() 80 | 81 | pf_latent.close() 82 | pf_pos.close() 83 | pf_measurement.close() 84 | 85 | 86 | 87 | -------------------------------------------------------------------------------- /data_processing/fitting/fit_latent.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the experimental code for paper ``Ma et al., OpenSVBRDF: A Database of Measured Spatially-Varying Reflectance, SIGGRAPH Asia 2023``. 3 | This script is suboptimal and experimental. 4 | There may be redundant lines and functionalities. 5 | 6 | Xiaohe Ma, 2024/02 7 | ''' 8 | 9 | 10 | ''' 11 | This script performs the fitting via differentiable rendering. 12 | 13 | ''' 14 | 15 | import os 16 | os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1" 17 | import argparse 18 | from pathlib import Path 19 | import torch.utils.data as data 20 | import sys 21 | 22 | from latent_controller import LatentController 23 | from latent_solver import LatentModelSolver 24 | from latent_dataset import LatentDataset 25 | from latent_mlp import LatentMLP 26 | from utils import setup_seed, setup_multiprocess 27 | 28 | 29 | TORCH_RENDER_PATH = "../../torch_renderer/" 30 | sys.path.append(TORCH_RENDER_PATH) 31 | from torch_render import TorchRender 32 | from setup_config import SetupConfig 33 | 34 | PROJECT_ROOT = Path(__file__).parent.absolute() 35 | 36 | train_device = "cuda:0" 37 | 38 | 39 | OBJECT_CONFIG = { 40 | "satin": { 41 | "axay_range": 0.8, "ps_range": 10, "lambda_axay": 0.01, "lambda_m": 0.05 42 | }, 43 | "fabric": { 44 | "axay_range": 0.5, "ps_range": 10, "lambda_axay": 0.1, "lambda_m": 0.001 45 | }, 46 | "leather": { 47 | "axay_range": 0.5, "ps_range": 10, "lambda_axay": 0.01, "lambda_m": 0.005 48 | }, 49 | "paper": { 50 | "axay_range": 0.5, "ps_range": 20, "lambda_axay": 0.05, "lambda_m": 0.005 51 | }, 52 | "wallpaper": { 53 | "axay_range": 0.5, "ps_range": 10, "lambda_axay": 0.1, "lambda_m": 0.005 54 | }, 55 | "wood": { 56 | "axay_range": 1, "ps_range": 10, "lambda_axay": 0.01, "lambda_m": 0.005 57 | }, 58 | "woven": { 59 | "axay_range": 1, "ps_range": 10, "lambda_axay": 0.01, "lambda_m": 0.005 60 | }, 61 | "metal": { 62 | "axay_range": 0.8, "ps_range": 10, "lambda_axay": 0.05, "lambda_m": 0.005 63 | }, 64 | "ceramic": { 65 | "axay_range": 0.5, "ps_range": 10, "lambda_axay": 0.1, "lambda_m": 0.005 66 | } 67 | } 68 | 69 | 70 | 71 | def parse_args(): 72 | parser = argparse.ArgumentParser() 73 | parser.add_argument("sample_class",type=str) 74 | parser.add_argument("data_root",type=str) 75 | parser.add_argument("save_root",default="texture_maps/") 76 | 77 | parser.add_argument('--train_device', type=str, default="cuda:0") 78 | parser.add_argument('--iter', type=int, default=120000) 79 | parser.add_argument("--main_cam_id",type=int,default=0) 80 | parser.add_argument("--tex_resolution",type=int,default=1024) 81 | 82 | parser.add_argument('--config_dir',type=str, default="../../torch_renderer/wallet_of_torch_renderer/lightstage/") 83 | parser.add_argument('--model_file',type=str, default="../../model/latent_48_24_500000_2437925.pkl") 84 | 85 | args, _ = parser.parse_known_args() 86 | 87 | global train_device 88 | train_device = args.train_device 89 | return args 90 | 91 | 92 | def main(): 93 | setup_seed(20) 94 | setup_multiprocess() 95 | 96 | args = parse_args() 97 | os.makedirs(args.save_root, exist_ok=True) 98 | 99 | setup = SetupConfig(args.config_dir, low_res=False) 100 | torch_render = TorchRender(setup) 101 | 102 | latent_controller = LatentController(args.model_file, 103 | args.config_dir, 104 | train_device) 105 | batch_size = 2**7 106 | 107 | latent_dataset = LatentDataset(latent_controller, args.save_root,args.data_root,batch_size, train_device) 108 | latent_loader = data.DataLoader(latent_dataset, batch_size=None, num_workers=0) 109 | 110 | ggx_mlp_config = { 111 | "position_texture" : args.save_root+"pos_texture.exr", 112 | "texture_resolution" : args.tex_resolution, 113 | "ps_range" :OBJECT_CONFIG[args.sample_class]["ps_range"], 114 | "axay_range" : OBJECT_CONFIG[args.sample_class]["axay_range"], 115 | "lambda_axay" :OBJECT_CONFIG[args.sample_class]["lambda_axay"], 116 | "lambda_m" : OBJECT_CONFIG[args.sample_class]["lambda_axay"] 117 | } 118 | 119 | ggx_mlp = LatentMLP(torch_render, train_device, **ggx_mlp_config).to(train_device) 120 | 121 | training_args = { 122 | "lr": 1e-3, 123 | "num_iters": args.iter, 124 | "device": train_device 125 | } 126 | 127 | solver = LatentModelSolver(ggx_mlp, latent_loader, **training_args) 128 | solver.train(latent_controller, args.data_root, args.save_root) 129 | 130 | 131 | if __name__ == "__main__": 132 | main() -------------------------------------------------------------------------------- /data_processing/fitting/latent_controller.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import os.path as osp 4 | import sys 5 | 6 | TORCH_RENDER_PATH = "../torch_renderer/" 7 | sys.path.append(TORCH_RENDER_PATH) 8 | 9 | from setup_config import SetupConfig 10 | 11 | AUTOENCODER_PATH="../finetune/" 12 | sys.path.append(AUTOENCODER_PATH) 13 | from AUTO_planar_scanner_net_inference import PlanarScannerNet 14 | 15 | 16 | class LatentController(object): 17 | ''' 18 | LatentController includes methods for loading latent models, measurements, and latent data. 19 | 20 | Methods: 21 | __init__: Initializes the LatentController with the model file, setup configuration directory, and device. 22 | load_model: Loads a trained model from the model file and updates the parameters to PlanarScannerNet. 23 | load_measurements: Loads the measurement data. 24 | load_latent_data: Loads the latent data from the specified directory. 25 | pred: Makes predictions based on the input latent data of color and shape. 26 | get_light_pattern: Retrieves the light pattern. 27 | 28 | ''' 29 | def __init__( 30 | self, 31 | model_file: str, 32 | setup_config_dir: str, 33 | device: str, 34 | ) -> None: 35 | self.setup = SetupConfig(setup_config_dir) 36 | self.model = self.load_model(model_file, device) 37 | self.device = device 38 | self.tex_resolution = 1024 39 | self.model_root = osp.dirname(model_file) 40 | 41 | def load_model( 42 | self, 43 | model_file: str, 44 | device: str, 45 | ) -> PlanarScannerNet: 46 | """ 47 | Load latent model 48 | """ 49 | train_configs = {} 50 | train_configs["training_device"] = 0 51 | train_configs["lumitexel_length"] = 64*64*3 52 | train_configs["shape_latent_len"] = 48 53 | train_configs["color_latent_len"] = 8 54 | inference_net = PlanarScannerNet(train_configs) 55 | pretrained_dict = torch.load(model_file) 56 | print("loading trained model...") 57 | 58 | something_not_found = False 59 | model_dict = inference_net.state_dict() 60 | for k,_ in model_dict.items(): 61 | if k not in pretrained_dict: 62 | print("not found:", k) 63 | something_not_found = True 64 | if something_not_found: 65 | exit() 66 | 67 | model_dict = inference_net.state_dict() 68 | pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} 69 | model_dict.update(pretrained_dict) 70 | inference_net.load_state_dict(model_dict) 71 | for p in inference_net.parameters(): 72 | p.requires_grad=False 73 | inference_net.to(device) 74 | inference_net.eval() 75 | return inference_net 76 | 77 | def load_measurements( 78 | self, 79 | data_root: str, 80 | ) -> torch.Tensor: 81 | m_path = osp.join(data_root, f"line_measurements_{self.tex_resolution}.bin") 82 | if osp.exists(m_path): 83 | measurements = np.fromfile(m_path, np.float32).reshape(-1, 64, 3) 84 | measurements = torch.from_numpy(measurements) 85 | else: 86 | measurements = None 87 | return measurements 88 | 89 | def load_latent_data( 90 | self, 91 | data_root: str, 92 | ): 93 | shape_latent_len = self.model.lumitexel_net.shape_latent_len 94 | color_latent_len = self.model.lumitexel_net.color_latent_len 95 | all_latent_len = shape_latent_len + 3 * color_latent_len 96 | 97 | pf_latent = open(osp.join(data_root, f"latent/pass3_latent_{self.tex_resolution}.bin"), "rb") 98 | pf_latent.seek(0, 2) 99 | texel_num = pf_latent.tell() // all_latent_len // 4 100 | pf_latent.seek(0, 0) 101 | print("texel num : ", texel_num) 102 | positions = np.fromfile(osp.join(data_root, f"texture_1024/positions.bin"), np.float32).reshape([-1, 3]) 103 | positions = torch.from_numpy(positions) 104 | assert positions.size(0) == texel_num 105 | 106 | latent_code = np.fromfile(pf_latent, np.float32).reshape([texel_num, all_latent_len]) 107 | latent_code = torch.from_numpy(latent_code) 108 | 109 | color_latent = latent_code[:, :3 * color_latent_len].reshape([texel_num, 3, color_latent_len]) 110 | shape_latent = latent_code[:, 3 * color_latent_len:].reshape([texel_num, 1, shape_latent_len]).repeat(1, 3, 1) 111 | 112 | color_shape_latent = torch.cat([color_latent, shape_latent],dim=-1) 113 | 114 | self.color_shape_latent = color_shape_latent.to(self.device) 115 | self.positions = positions.to(self.device) 116 | return color_shape_latent, positions 117 | 118 | def pred( 119 | self, 120 | color_shape_latent 121 | ): 122 | batch_size = color_shape_latent.size(0) 123 | latent_len = color_shape_latent.size(2) 124 | tmp_latent = color_shape_latent.reshape([batch_size*3, latent_len]) 125 | 126 | _, tmp_nn_lumi = self.model(tmp_latent, input_is_latent=True) 127 | tmp_nn_lumi = torch.max(torch.zeros_like(tmp_nn_lumi), tmp_nn_lumi) 128 | tmp_nn_lumi = tmp_nn_lumi.reshape([batch_size, 3, -1]).permute(0, 2, 1) 129 | return tmp_nn_lumi 130 | 131 | def get_light_pattern(self) -> torch.Tensor: 132 | path = f"{self.model_root}/opt_W_64.bin" 133 | lp = None 134 | if osp.exists(path): 135 | lp = torch.from_numpy(np.fromfile(path, np.float32).reshape(64, -1, 3)) 136 | lp = lp.permute(1, 2, 0) 137 | return lp -------------------------------------------------------------------------------- /data_processing/fitting/latent_dataset.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.utils.data as data 3 | from pathlib import Path 4 | from utils import write_rgb_image 5 | from latent_controller import LatentController 6 | 7 | 8 | class LatentDataset(data.Dataset): 9 | ''' 10 | LatentDataset builds a dataset from the model file for fitting. 11 | 12 | Methods: 13 | __init__: Initializes the dataset with the latent controller, save root, data root, batch size, and device. 14 | valid_test: Removes invalid points from the dataset. 15 | gen_uv: Generates UV coordinates for further texture mapping. 16 | __load_fitting_data: Loads the fitting data from the specified directories. 17 | __len__: Returns the total number of valid texels in the dataset. 18 | __getitem__: Retrieves a batch of data at the specified index. 19 | get_light_pattern: Retrieves the light pattern from the latent controller. 20 | 21 | ''' 22 | def __init__( 23 | self, 24 | latent_controller: LatentController, 25 | save_root: str, 26 | data_root: str, 27 | batch_size: int, 28 | device: str, 29 | ) -> None: 30 | super().__init__() 31 | self.save_root = Path(save_root) 32 | self.latent_data_root = Path(data_root+"latent/") 33 | self.data_root = Path(data_root) 34 | self.m_data_root = Path(data_root+"texture_1024/") 35 | self.latent_controller = latent_controller 36 | self.batch_size = batch_size 37 | self.device = device 38 | 39 | latent_code, point_pos, uvs = self.__load_fitting_data() 40 | self.latent_code = latent_code 41 | self.point_pos = point_pos 42 | self.uvs = uvs 43 | self.pixel_num = self.latent_code.size(0) 44 | self.measurements = self.latent_controller.load_measurements(str(self.m_data_root)) 45 | 46 | self.valid_test() 47 | 48 | def valid_test(self): 49 | print(f"Total texel num: {self.latent_code.size(0)}") 50 | latent_sum = torch.sum(self.latent_code, dim=(1, 2)) 51 | valid_mask = ~torch.isnan(latent_sum) 52 | self.latent_code = self.latent_code[valid_mask] 53 | self.point_pos = self.point_pos[valid_mask] 54 | self.uvs = self.uvs[valid_mask] 55 | self.measurements = self.measurements[valid_mask] 56 | self.pixel_num = self.latent_code.size(0) 57 | print(f"Valid texel num: {self.latent_code.size(0)}") 58 | 59 | def gen_uv(self, resolution=1024): 60 | half_dx = 0.5 / resolution 61 | half_dy = 0.5 / resolution 62 | xs = torch.linspace(half_dx, 1-half_dx, resolution, device=self.device) 63 | ys = torch.linspace(half_dx, 1-half_dy, resolution, device=self.device) 64 | xv, yv = torch.meshgrid([1 - xs, ys]) 65 | xy = torch.stack((yv.flatten(), xv.flatten())).t() 66 | return xy 67 | 68 | def __load_fitting_data(self): 69 | latent_code, positions = self.latent_controller.load_latent_data(str(self.data_root)) 70 | 71 | uv_map = self.gen_uv() 72 | 73 | # save position texture 74 | pos_texture = positions.view(1024, 1024, 3).cpu().numpy() 75 | path = str(self.save_root / "pos_texture.exr") 76 | write_rgb_image(path, pos_texture) 77 | print(f"Generate position texture: {path}") 78 | 79 | return latent_code, positions, uv_map 80 | 81 | def __len__(self) -> int: 82 | return self.pixel_num 83 | 84 | def __getitem__(self, index: int): 85 | 86 | batch = torch.rand((self.batch_size, ), device=self.device, dtype=torch.float32) 87 | batch = (batch * self.pixel_num).long().cpu() 88 | latent_code = self.latent_code[batch] 89 | point_pos = self.point_pos[batch] 90 | uv = self.uvs[batch] 91 | 92 | m = self.measurements[batch] if self.measurements is not None else None 93 | lumitexel = self.latent_controller.pred(latent_code.to(self.device)) 94 | 95 | train_data = { 96 | 'lumitexel': lumitexel, 97 | 'position': point_pos, 98 | 'uv': uv, 99 | 'measurements': m 100 | } 101 | return train_data 102 | 103 | def get_light_pattern(self): 104 | return self.latent_controller.get_light_pattern() -------------------------------------------------------------------------------- /data_processing/fitting/latent_mlp.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import os.path as osp 4 | from tqdm import tqdm 5 | import math 6 | import sys 7 | 8 | from texture import Texture 9 | from mlp_model import MLPModel 10 | 11 | TORCH_RENDER_PATH = "../torch_renderer/" 12 | sys.path.append(TORCH_RENDER_PATH) 13 | from torch_render import TorchRender 14 | from utils import write_rgb_image, read_exr 15 | 16 | 17 | class LatentMLP(nn.Module): 18 | ''' 19 | LatentMLP defines a multilayer perceptron model. 20 | 21 | Methods: 22 | pred_params: Predicts parameters such as the lighting direction and roughness from the input UV coordinates. 23 | export_texture: Saves the final texture results as .exr and .png files. 24 | forward: Defines forward propagation. 25 | loss: Calculates the loss during training. 26 | 27 | ''' 28 | def __init__( 29 | self, 30 | torch_render: TorchRender, 31 | device: str, 32 | **kwargs 33 | ) -> None: 34 | super().__init__() 35 | self.torch_render = torch_render 36 | self.device = device 37 | self.step = 0 38 | 39 | self.pos_texture_file = kwargs.pop("position_texture", None) 40 | self.resolution = kwargs.pop("texture_resolution", 512) 41 | self.ps_range = kwargs.pop("ps_range", 5) 42 | self.axay_range = kwargs.pop("axay_range", 0.497) 43 | self.lambda_axay = kwargs.pop("lambda_axay", 0.1) 44 | self.lambda_m = kwargs.pop("lambda_m", 0.01) 45 | 46 | if self.pos_texture_file is None: 47 | print("Error, position_texture must be given.") 48 | 49 | self.inner_dim = 24 50 | self.param_encoding = Texture(self.resolution, self.resolution, self.inner_dim) 51 | 52 | pos_tex = torch.from_numpy(read_exr(self.pos_texture_file)) 53 | for i in range(pos_tex.size(2)): 54 | pos_tex[:, :, i] = torch.Tensor.flipud(pos_tex[:, :, i]) 55 | self.pos_texture = Texture(*(pos_tex.shape)) 56 | self.pos_texture.set_parameters(pos_tex, False) 57 | 58 | self.n2d_mlp = MLPModel([self.inner_dim, 128, 128, 128, 2], normalizaiton=None, output_activation="sigmoid") 59 | self.theta_mlp = MLPModel([self.inner_dim, 128, 128, 128, 1], normalizaiton=None, output_activation="sigmoid") 60 | self.diffuse_albedo_mlp = MLPModel([self.inner_dim, 128, 128, 128, 3], normalizaiton=None, output_activation="sigmoid") 61 | self.specular_albedo_mlp = MLPModel([self.inner_dim, 128, 128, 128, 3], normalizaiton=None, output_activation="sigmoid") 62 | self.roughness_mlp = MLPModel([self.inner_dim, 128, 128, 128, 2], normalizaiton=None, output_activation="sigmoid") 63 | 64 | self.albedo_scalar = nn.Parameter(torch.ones(1), requires_grad=False) 65 | 66 | def pred_params(self, uvs): 67 | neck = self.param_encoding(uvs) 68 | 69 | n2d = self.n2d_mlp(neck) 70 | theta = self.theta_mlp(neck) * torch.pi * 2 71 | roughness = self.roughness_mlp(neck) * self.axay_range + 0.006 72 | 73 | diffuse_albedo = self.diffuse_albedo_mlp(neck) * self.albedo_scalar 74 | specular_albedo = self.specular_albedo_mlp(neck) * self.ps_range * self.albedo_scalar 75 | 76 | input_params = torch.cat([n2d, theta, roughness, diffuse_albedo, specular_albedo], dim=1) 77 | return input_params 78 | 79 | def get_position(self, uvs): 80 | p = self.pos_texture(uvs.detach()) 81 | return p 82 | 83 | def forward( 84 | self, 85 | uvs: torch.Tensor, 86 | scalar: float = 625, 87 | ): 88 | pred_params = self.pred_params(uvs) 89 | point_pos = self.get_position(uvs) 90 | 91 | _, end_points = self.torch_render.generate_lumitexel( 92 | pred_params, 93 | point_pos, 94 | pd_ps_wanted="both", 95 | ) 96 | 97 | diffuse_lumi = end_points["diff_lumi"] * scalar 98 | specular_lumi = end_points["spec_lumi"] * scalar 99 | 100 | end_points['pred_params'] = pred_params 101 | end_points['point_pos'] = point_pos 102 | 103 | return diffuse_lumi, specular_lumi, end_points 104 | 105 | def export_texture(self, exr_dir=None, png_dir=None, resolution=1024): 106 | half_dx = 0.5 / resolution 107 | half_dy = 0.5 / resolution 108 | xs = torch.linspace(half_dx, 1-half_dx, resolution, device=self.device) 109 | ys = torch.linspace(half_dx, 1-half_dy, resolution, device=self.device) 110 | xv, yv = torch.meshgrid([1 - xs, ys], indexing="ij") 111 | 112 | xy = torch.stack((yv.flatten(), xv.flatten())).t() 113 | 114 | with torch.no_grad(): 115 | batch_size = 1024 116 | pred_params, pred_pos, pred_normal, pred_tangent = [], [], [], [] 117 | for i in tqdm(range(xy.size(0) // batch_size + 1)): 118 | start = i * batch_size 119 | end = min((i + 1) * batch_size, xy.size(0)) 120 | if start == end: 121 | continue 122 | xy_tmp = xy[start:end] 123 | _, _, end_points = self.forward(xy_tmp) 124 | pred_params.append(end_points["pred_params"]) 125 | pred_pos.append(end_points["point_pos"]) 126 | pred_normal.append(end_points["n"]) 127 | pred_tangent.append(end_points['t']) 128 | 129 | pred_params = torch.cat(pred_params, dim=0).to(self.device) 130 | pred_pos = torch.cat(pred_pos, dim=0).to(self.device) 131 | pred_normal = torch.cat(pred_normal, dim=0).to(self.device) 132 | pred_tangent = torch.cat(pred_tangent, dim=0).to(self.device) 133 | 134 | # process ax, ay and tangent 135 | n, t = pred_normal, pred_tangent 136 | b = torch.cross(n, t) 137 | pred_tangent = torch.where(pred_params[:, [3]] > pred_params[:, [4]], pred_tangent, b) 138 | pred_params[:, [3, 4]] = torch.where(pred_params[:, [3]] > pred_params[:, [4]], pred_params[:, [3, 4]], pred_params[:, [4, 3]]) 139 | 140 | pred_params[:, [2]] = torch.where(pred_params[:, [3]] > pred_params[:, [4]], pred_params[:, [2]], pred_params[:, [2]] + torch.pi) 141 | pred_params[:, [2]] = torch.where(pred_params[:, [2]] <= torch.pi * 2, pred_params[:, [2]], pred_params[:, [2]] - torch.pi * 2) 142 | 143 | params_texture = pred_params.view(resolution, resolution, -1).detach().cpu() 144 | theta_texture = params_texture[:, :, [2]].repeat(1, 1, 3) 145 | pd_texture = params_texture[:, :, 5:8] / self.albedo_scalar.detach().cpu() 146 | ps_texture = params_texture[:, :, 8:11] / self.albedo_scalar.detach().cpu() 147 | ax_ay_img = torch.zeros_like(params_texture[:, :, :3]) 148 | ax_ay_img[:, :, :2] = params_texture[:, :, [3, 4]] 149 | ax_ay_texture = ax_ay_img 150 | 151 | n2d_texture = torch.zeros_like(params_texture[:, :, :3]) 152 | n2d_texture[:, :, :2] = params_texture[:, :, [0, 1]] 153 | 154 | normal_texture = pred_normal.view(resolution, resolution, -1).detach().cpu() 155 | normal_texture = normal_texture * 0.5 + 0.5 156 | 157 | tangent_texture = pred_tangent.view(resolution, resolution, -1).detach().cpu() 158 | tangent_texture = tangent_texture * 0.5 + 0.5 159 | 160 | if exr_dir is not None: 161 | write_rgb_image(osp.join(exr_dir, "ax_ay_texture.exr"), ax_ay_texture.numpy()) 162 | write_rgb_image(osp.join(exr_dir, "pd_texture.exr"), pd_texture.numpy()) 163 | write_rgb_image(osp.join(exr_dir, "ps_texture.exr"), ps_texture.numpy()) 164 | write_rgb_image(osp.join(exr_dir, "normal_texture.exr"), normal_texture.numpy()) 165 | write_rgb_image(osp.join(exr_dir, "tangent_texture.exr"), tangent_texture.numpy()) 166 | write_rgb_image(osp.join(exr_dir, "theta_texture.exr"), theta_texture.numpy()) 167 | write_rgb_image(osp.join(exr_dir, "n2d_texture.exr"), n2d_texture.numpy()) 168 | 169 | if png_dir is not None: 170 | write_rgb_image(osp.join(png_dir, "ax_ay.png"), ax_ay_texture.numpy() ** (1 / 2.2) * 255) 171 | write_rgb_image(osp.join(png_dir, "pd.png"), pd_texture.numpy() ** (1 / 2.2) * 255) 172 | write_rgb_image(osp.join(png_dir, "ps.png"), (ps_texture.numpy() / 10) ** (1 / 2.2) * 255) 173 | write_rgb_image(osp.join(png_dir, "normal.png"), normal_texture.numpy() * 255) 174 | write_rgb_image(osp.join(png_dir, "tangent.png"), tangent_texture.numpy() * 255) 175 | write_rgb_image(osp.join(png_dir, "theta.png"), theta_texture.numpy() / math.pi / 2 * 255) 176 | write_rgb_image(osp.join(png_dir, "n2d.png"), n2d_texture.numpy() * 255) 177 | 178 | 179 | textures = torch.stack([theta_texture / torch.pi / 2, ax_ay_texture, pd_texture, 180 | ps_texture / 5, normal_texture, tangent_texture], dim=0) 181 | textures = textures.permute(0, 3, 1, 2) 182 | return textures 183 | 184 | def loss( 185 | self, 186 | uvs: torch.Tensor, 187 | lumi: torch.Tensor, 188 | lp: torch.Tensor = None, 189 | m: torch.Tensor = None, 190 | ): 191 | pred_diff_lumi, pred_spec_lumi, end_points = self.forward(uvs, 3e3 / math.pi) 192 | pred_lumi = pred_diff_lumi + pred_spec_lumi 193 | 194 | if self.step == 1: 195 | print(f"Gt Lumitexel: max_val {lumi.max()}, mean_val: {lumi.mean()}") 196 | print(f"Pred Lumitexel: max_val {pred_lumi.max()}, mean_val: {pred_lumi.detach().mean()}") 197 | s = pred_lumi.detach().mean() / lumi.mean() 198 | self.albedo_scalar.data[0] = s 199 | print(f"Set albedo_scalar to {s}") 200 | 201 | latent_radiance = lumi.unsqueeze(3) * lp.unsqueeze(0) 202 | latent_radiance = torch.sum(latent_radiance, dim=1).transpose(2, 1) 203 | loss_m = torch.nn.MSELoss()(latent_radiance * self.albedo_scalar, m * self.albedo_scalar) 204 | print(f"Latent measurements MSE: {loss_m.item()}") 205 | 206 | # Loss lumitexel 207 | loss_lumi = torch.nn.MSELoss()(pred_lumi, lumi * self.albedo_scalar) 208 | 209 | # Loss: ax, ay regularization 210 | ax = end_points['pred_params'][:, [3]] 211 | ay = end_points['pred_params'][:, [4]] 212 | loss_ax_ay = torch.mean(ax * ay) 213 | 214 | # Loss measurements 215 | radiance = pred_lumi.unsqueeze(3) * lp.unsqueeze(0) 216 | radiance = torch.sum(radiance, dim=1).transpose(2, 1) 217 | loss_m = torch.nn.MSELoss()(radiance, m * self.albedo_scalar) 218 | 219 | loss = loss_lumi + loss_ax_ay * self.lambda_axay + loss_m * self.lambda_m 220 | 221 | meta = { 222 | 'loss_lumi': loss_lumi.item(), 223 | 'loss_ax_ay': loss_ax_ay.item(), 224 | 'loss_m': loss_m.item(), 225 | 'pred_lumi': pred_lumi.detach(), 226 | } 227 | return loss, meta 228 | -------------------------------------------------------------------------------- /data_processing/fitting/latent_solver.py: -------------------------------------------------------------------------------- 1 | import logging 2 | from pathlib import Path 3 | import torch 4 | import torch.utils.data as data 5 | 6 | from latent_mlp import LatentMLP 7 | from latent_controller import LatentController 8 | 9 | from torch_render import TorchRender 10 | 11 | 12 | class LatentModelSolver(object): 13 | ''' 14 | LatentModelSolver is a model solver used to train and optimize lighting parameters with the Adam optimizer. 15 | 16 | Methods: 17 | __init__: Initializes the LatentModelSolver with a model, data loader, and various parameters. 18 | train: Performs fitting, calculates losses, and saves final texture results. 19 | ''' 20 | def __init__( 21 | self, 22 | model: LatentMLP, 23 | data_loader: data.DataLoader, 24 | **kwargs 25 | ) -> None: 26 | self.model = model 27 | self.data_loader = data_loader 28 | 29 | self.lr = kwargs.pop("lr", 0.0005) 30 | self.num_iters = kwargs.pop("num_iters", 1000000) 31 | self.device = kwargs.pop("device", "cuda:0") 32 | 33 | self.step = 0 # current training step 34 | self.model.to(self.device) 35 | self.optimizer = torch.optim.Adam( 36 | model.parameters(), 37 | lr=self.lr, 38 | ) 39 | 40 | def train(self, latent_controller: LatentController, fitting_dir: str, output_dir: str): 41 | 42 | torch_render = TorchRender(latent_controller.setup) 43 | lp = self.data_loader.dataset.get_light_pattern().to(self.device) 44 | 45 | for train_data in self.data_loader: 46 | 47 | self.model.train() 48 | 49 | self.step += 1 50 | self.model.step += 1 51 | assert(self.step == self.model.step) 52 | if self.step > self.num_iters: 53 | break 54 | 55 | batch_uvs = train_data['uv'].to(self.device) 56 | batch_lumi = train_data['lumitexel'].to(self.device) 57 | batch_m = train_data['measurements'].to(self.device) 58 | loss, meta = self.model.loss(batch_uvs, batch_lumi, lp, batch_m) 59 | 60 | self.optimizer.zero_grad() 61 | loss.backward() 62 | self.optimizer.step() 63 | if self.step % 100 == 0: 64 | log_info = "[%d] Loss: %.6f" % (self.step, loss.item()) 65 | for key in meta: 66 | if "loss" in key: 67 | log_info = log_info + " %s: %.6f" % (key, meta[key]) 68 | print(log_info) 69 | 70 | if output_dir is not None : 71 | output_dir = Path(output_dir) 72 | # export texture map 73 | self.model.eval() 74 | textures = self.model.export_texture(str(output_dir), str(output_dir)) 75 | 76 | 77 | -------------------------------------------------------------------------------- /data_processing/fitting/mlp_model.py: -------------------------------------------------------------------------------- 1 | from typing import List 2 | import torch.nn as nn 3 | from collections import OrderedDict 4 | import torch.nn.init as init 5 | import math 6 | 7 | 8 | class MLPModel(nn.Module): 9 | ''' 10 | MLPModel defines a simple multilayer perceptron model. 11 | 12 | Methods: 13 | __init__: Initializes the model with the input dimension list. Linear layers are added. Normalization layers and activation functions are included if specified. 14 | forward: Defines forward propagation. 15 | reset_parameters: Resets the parameters of the linear layers using Kaiming initialization. 16 | ''' 17 | 18 | def __init__( 19 | self, 20 | dim_list: List[int], 21 | normalizaiton='BN', 22 | activation='leaky_relu', 23 | output_activation=None, 24 | ): 25 | super(MLPModel, self).__init__() 26 | 27 | layers = OrderedDict() 28 | for i in range(len(dim_list)-2): 29 | layers[f"Linear_{i}"] = nn.Linear(dim_list[i], dim_list[i+1]) 30 | if normalizaiton == "BN": 31 | layers["Batchnorm_{}".format(i)] = nn.BatchNorm1d(dim_list[i+1]) 32 | elif normalizaiton == "LN": 33 | layers["Layernorm_{}".format(i)] = nn.LayerNorm(dim_list[i+1]) 34 | if activation == 'leaky_relu': 35 | layers[f"LeakyRelu_{i}"] = nn.LeakyReLU(0.2, inplace=True) 36 | elif activation == 'relu': 37 | layers[f"Relu_{i}"] = nn.ReLU() 38 | 39 | i = len(dim_list)-2 40 | layers[f"Linear_{i}"] = nn.Linear(dim_list[i], dim_list[i+1]) 41 | 42 | if output_activation is not None and output_activation == "sigmoid": 43 | layers["Sigmoid"] = nn.Sigmoid() 44 | 45 | self.decode_part = nn.Sequential(layers) 46 | 47 | def forward(self, x): 48 | out = self.decode_part(x) 49 | return out 50 | 51 | def reset_parameters(self) -> None: 52 | 53 | def reset(m): 54 | if type(m) == nn.Linear: 55 | init.kaiming_uniform_(m.weight, a=math.sqrt(5)) 56 | if m.bias is not None: 57 | fan_in, _ = init._calculate_fan_in_and_fan_out(m.weight) 58 | bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 59 | init.uniform_(m.bias, -bound, bound) 60 | 61 | self.decode_part.apply(reset) 62 | -------------------------------------------------------------------------------- /data_processing/fitting/run.sh: -------------------------------------------------------------------------------- 1 | for variable in 18 2 | do 3 | class_name=paper 4 | formatted_variable=$(printf "%04d" $variable) 5 | name=$class_name$formatted_variable 6 | tex_resolution=1024 7 | 8 | data_root=../../database_data/$name/output/ 9 | save_root=../../database_data/$name/output/texture_maps/ 10 | config_dir=../torch_renderer/wallet_of_torch_renderer/lightstage/ 11 | model_file=../../database_model/latent_48_24_500000_2437925.pkl 12 | train_device="cuda:0" 13 | python fit_latent.py $class_name $data_root $save_root --train_device $train_device --config_dir $config_dir --model_file $model_file 14 | 15 | done 16 | -------------------------------------------------------------------------------- /data_processing/fitting/texture.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import torch.nn as nn 4 | from torch.nn import init 5 | import torch.nn.functional as F 6 | 7 | 8 | class Texture(nn.Module): 9 | ''' 10 | Texture implements simple texture mapping. 11 | 12 | Methods: 13 | forward: Defines forward propagation. Performs texture mapping of the input UV coordinates. 14 | reset_parameters: Resets the parameters using Xavier initialization. 15 | set_parameters: Sets the texture map parameters. 16 | ''' 17 | 18 | def __init__(self, width, height, feature_num): 19 | super(Texture, self).__init__() 20 | self.width = width 21 | self.height = height 22 | self.feature_num = feature_num 23 | self.params = nn.Parameter(torch.Tensor(1, feature_num, width, height)) 24 | self.reset_parameters() 25 | 26 | def reset_parameters(self): 27 | init.xavier_uniform_(self.params.data) 28 | 29 | def set_parameters(self, params, requires_grad): 30 | param_data = params.permute(2, 0, 1).unsqueeze(0).to(self.params.device) 31 | self.params.data.copy_(param_data) 32 | self.params.requires_grad = requires_grad 33 | 34 | def forward(self, uv_): 35 | batch = uv_.size(0) 36 | uv = uv_ * 2.0 - 1.0 37 | 38 | uv = uv.view(1, -1, 1, 2) 39 | y = F.grid_sample(self.params, uv, align_corners=True) 40 | 41 | y = y.view(-1, batch).transpose(1, 0) 42 | return y -------------------------------------------------------------------------------- /data_processing/fitting/utils.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import torch 4 | import sys 5 | def read_exr(filename: str) -> np.ndarray: 6 | """ 7 | Read exr image 8 | 9 | Args: 10 | filename: image filename 11 | 12 | Returns: 13 | image: a ndarray of shape (H, W, C). The three channels are in 14 | RBG format. 15 | """ 16 | assert(filename.endswith(".exr")) 17 | image = cv2.imread(filename, cv2.IMREAD_ANYDEPTH | cv2.IMREAD_ANYCOLOR) 18 | if len(image.shape) == 3: 19 | image = image[:, :, [2, 1, 0]] 20 | return image 21 | 22 | 23 | def read_ldr(filename: str) -> np.array: 24 | """ 25 | Returns: 26 | image: RGB format 27 | """ 28 | image = cv2.imread(filename) 29 | if image is None: 30 | return image 31 | image = image[:, :, ::-1] 32 | return image 33 | 34 | 35 | def write_rgb_image(filename: str, image: np.array): 36 | """ 37 | OpenCV!! 38 | 39 | Args: 40 | image: RGB image (H, W, 3) 41 | """ 42 | image = image.astype(np.float32) 43 | 44 | if len(image.shape) == 3: 45 | assert(image.shape[2] in [1, 3]) 46 | # convert RGB to BRG 47 | image = image[:, :, ::-1] 48 | 49 | cv2.imwrite(filename, image) 50 | 51 | def setup_seed(seed): 52 | torch.manual_seed(seed) 53 | torch.cuda.manual_seed_all(seed) 54 | np.random.seed(seed) 55 | torch.backends.cudnn.deterministic = True 56 | 57 | 58 | def setup_multiprocess(): 59 | if sys.platform == 'linux': 60 | torch.multiprocessing.set_start_method('spawn', force=True) -------------------------------------------------------------------------------- /data_processing/generate_uv/extract_plane.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import argparse 4 | os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1" 5 | import cv2 6 | import sys 7 | sys.path.append("../camera_related/") 8 | from camera_config import Camera 9 | 10 | def undistort_masks(cameras, ldr_path, hdr_path, down_size): 11 | """ 12 | Undistort the mask and image files using the provided camera parameters. 13 | 14 | Parameters: 15 | cameras (list): List of Camera objects for each camera. 16 | ldr_path (str): Path to the LDR images. 17 | hdr_path (str): Path to the HDR images. 18 | down_size (int): Factor by which to downsize the images. 19 | """ 20 | for which_cam in range(len(cameras)): 21 | # Read and undistort LDR images 22 | mask = cv2.imread(ldr_path + f"mask_cam{which_cam:02d}.png") 23 | mask = cameras[which_cam].undistort_2steps(mask) 24 | cv2.imwrite(ldr_path + f"mask_udt_cam{which_cam:02d}.png", mask) 25 | 26 | mask = cv2.resize(mask, (mask.shape[1]//down_size, mask.shape[0]//down_size), cv2.INTER_NEAREST) 27 | cv2.imwrite(ldr_path + f"mask_udt_cam{which_cam:02d}_{down_size}.png", mask) 28 | 29 | img = cv2.imread(ldr_path + f"{which_cam}_0.png") 30 | img = cameras[which_cam].undistort_2steps(img) 31 | cv2.imwrite(ldr_path + f"{which_cam}_0_udt.png", img) 32 | 33 | # Undistort and save HDR mask for the main camera (which_cam == 0) 34 | if which_cam == 0: 35 | mask = cv2.imread(hdr_path + f"mask_cam{which_cam:02d}.exr", 6) 36 | mask = cameras[which_cam].undistort_2steps(mask) 37 | cv2.imwrite(hdr_path + f"mask_udt_cam{which_cam:02d}.png",np.clip(mask*255, 0, 255).astype(np.uint8)) 38 | cv2.imwrite(hdr_path + f"mask_udt_cam{which_cam:02d}.exr", mask) 39 | mask = cv2.resize(mask, (mask.shape[1]//down_size, mask.shape[0]//down_size), cv2.INTER_NEAREST) 40 | cv2.imwrite(hdr_path + f"mask_udt_cam{which_cam:02d}_{down_size}.png", np.clip(mask*255, 0, 255).astype(np.uint8)) 41 | cv2.imwrite(hdr_path + f"mask_udt_cam{which_cam:02d}_{down_size}.exr", mask) 42 | return 43 | 44 | if __name__ == "__main__": 45 | parser = argparse.ArgumentParser(description='Extract measurements.') 46 | parser.add_argument('data_root', type=str) 47 | parser.add_argument('save_root', type=str) 48 | parser.add_argument('config_dir', type=str, help="Path to camera parameters.") 49 | parser.add_argument('--cam_num', type=int, default=2) 50 | parser.add_argument('--main_cam_id', type=int, default=0) 51 | parser.add_argument('--texture_resolution', type=int, default=1024) 52 | parser.add_argument('--down_size',type=int, default=2) 53 | 54 | args = parser.parse_args() 55 | 56 | # Define paths for LDR and HDR images 57 | ldr_path = args.data_root + "sfm/" 58 | hdr_path = args.data_root + "raw_images/" 59 | output_root = args.data_root + args.save_root + f"texture_{args.texture_resolution}/" 60 | 61 | os.makedirs(output_root, exist_ok=True) 62 | 63 | # Load camera configurations 64 | cameras = [Camera(args.config_dir + f"intrinsic{which_cam}.yml", args.config_dir + f"extrinsic{which_cam}.yml") for which_cam in range(args.cam_num)] 65 | 66 | undistort_masks(cameras, ldr_path, hdr_path, args.down_size) 67 | 68 | 69 | height = cameras[0].get_height() 70 | width = cameras[0].get_width() 71 | 72 | positions = np.fromfile(args.config_dir + f"positions_{args.down_size}.bin", np.float32).reshape([height//args.down_size, width//args.down_size,3]) 73 | 74 | mask = cv2.imread(hdr_path + f"mask_udt_cam{args.main_cam_id:02d}_{args.down_size}.exr", 6) 75 | camera = cameras[args.main_cam_id] 76 | 77 | # Find valid indices where the mask is valid 78 | valid_idxes = np.where(mask[:,:,0] == 1.0) 79 | # Extract valid positions 80 | valid_positions = positions[valid_idxes[1],valid_idxes[0]].reshape([-1,3]) 81 | 82 | # Project valid positions to image coordinates and downsize 83 | uv = camera.project(valid_positions.T) / args.down_size 84 | uv = np.maximum(uv,0) 85 | uv[:,0] = np.minimum(uv[:,0],width//args.down_size-1) 86 | uv[:,1] = np.minimum(uv[:,1],height//args.down_size-1) 87 | 88 | # Determine bounding box of valid positions 89 | x_min, x_max = np.min(valid_positions[:, 0]), np.max(valid_positions[:, 0]) 90 | y_min, y_max = np.min(valid_positions[:, 1]), np.max(valid_positions[:, 1]) 91 | x_range = x_max - x_min 92 | y_range = y_max - y_min 93 | 94 | # Scale and expand bounding box for safe region 95 | scale = 0.065 96 | expand = 50 // args.down_size 97 | z = -122 98 | 99 | x_start = int(x_min + scale * x_range) 100 | x_end = int(x_max - scale * x_range) 101 | y_start = int(y_min + scale * y_range) 102 | y_end = int(y_max - scale * y_range) 103 | length = np.minimum(x_end - x_start, y_end - y_start) 104 | 105 | uvs = [] 106 | # Project corners of bounding box to image coordinates 107 | for step_x in range(2): 108 | for step_y in range(2): 109 | x = x_start + step_x * length 110 | y = y_start + step_y * length 111 | 112 | uv = camera.project(np.array([x,y,z]).reshape([3,1])) / args.down_size 113 | uvs.append(uv) 114 | 115 | # Calculate minimum and maximum UV coordinates and apply expansion 116 | uvs = np.concatenate(uvs, axis=0) 117 | uv_min = np.min(uvs,axis=0).astype(np.int32) - expand 118 | uv_max = np.max(uvs,axis=0).astype(np.int32) + expand 119 | 120 | # Generate grid of positions for texture mapping 121 | x = np.linspace(x_start, x_end, args.texture_resolution) 122 | y = np.linspace(y_start, y_end, args.texture_resolution) 123 | X, Y = np.meshgrid(x, y) 124 | Z = np.full_like(X, z) 125 | positions = np.stack([X, Y, Z], axis=-1).reshape(-1, 3) 126 | positions.astype(np.float32).tofile(output_root+"positions.bin") 127 | 128 | # Generate UV coordinates for texture mapping 129 | u = np.linspace(0, args.texture_resolution - 1, args.texture_resolution) 130 | v = np.linspace(0, args.texture_resolution - 1, args.texture_resolution) 131 | U, V = np.meshgrid(u, v) 132 | uvs = np.stack([U, V], axis=-1).reshape(-1, 2) 133 | uvs.astype(np.int32).tofile(output_root+"texturemap_uv.bin") 134 | 135 | for which_cam in range(args.cam_num): 136 | # Project positions to UV coordinates for each camera 137 | uv = cameras[which_cam].project(positions.T) / args.down_size 138 | uv.astype(np.float32).tofile(output_root+f"uvs_cam{args.main_cam_id}.bin") 139 | 140 | # Compute region of interest (ROI) for the main camera 141 | if which_cam == 0: 142 | 143 | uv = uv.reshape([args.texture_resolution, args.texture_resolution, 2]).astype(np.int32) 144 | 145 | u_delta = np.max(np.abs(uv[:,1:,0] - uv[:,:args.texture_resolution-1,0])) 146 | v_delta = np.max(np.abs(uv[1:,:,1] - uv[:args.texture_resolution-1,:,1])) 147 | 148 | roi = np.array([uv_min[0], uv_min[1], uv_max[0], uv_max[1], int(np.ceil(u_delta/2+1))*2-1, int(np.ceil(v_delta/2+1))*2-1]) 149 | 150 | roi.astype(np.int32).tofile(output_root+"roi_{}.bin".format(args.down_size)) 151 | 152 | 153 | 154 | 155 | -------------------------------------------------------------------------------- /data_processing/generate_uv/run.sh: -------------------------------------------------------------------------------- 1 | for variable in 18 2 | do 3 | root=../../database_data/ 4 | class_name=paper 5 | 6 | formatted_variable=$(printf "%04d" $variable) 7 | data_root=$root$class_name$formatted_variable/ 8 | 9 | save_root=output/ 10 | config_dir=../device_configuration/ 11 | 12 | cam_num=2 13 | main_cam_id=0 14 | texture_resolution=1024 15 | down_size=2 16 | 17 | python extract_plane.py $data_root $save_root $config_dir --cam_num $cam_num --main_cam_id $main_cam_id --texture_resolution $texture_resolution --down_size $down_size 18 | 19 | python warp.py $data_root $save_root --main_cam_id $main_cam_id --texture_resolution $texture_resolution --down_size $down_size 20 | 21 | done 22 | -------------------------------------------------------------------------------- /data_processing/generate_uv/warp.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import argparse 3 | import os 4 | os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1" 5 | import cv2 6 | import cv2.aruco as aruco 7 | import sys 8 | 9 | if __name__ == "__main__": 10 | parser = argparse.ArgumentParser(description='Extract sample plane.') 11 | parser.add_argument('data_root', type=str) 12 | parser.add_argument('save_root', type=str) 13 | parser.add_argument('--main_cam_id', type=int, default=0) 14 | parser.add_argument('--texture_resolution', type=int, default=1024) 15 | parser.add_argument('--down_size',type=int, default=2) 16 | 17 | args = parser.parse_args() 18 | 19 | data_root = args.data_root + "sfm/" 20 | output_root = args.data_root + args.save_root + f"texture_{args.texture_resolution}/" 21 | 22 | # Load the region of interest (ROI) from the file 23 | roi = np.fromfile(output_root+"roi_{}.bin".format(args.down_size), np.int32).reshape([6,]) 24 | 25 | # Initialize the ArUco dictionary for marker detection 26 | aruco_dict = aruco.getPredefinedDictionary(aruco.DICT_6X6_1000) 27 | 28 | 29 | frame = [] 30 | gray = [] 31 | keys = [] 32 | descriptors = [] 33 | 34 | # Create a SIFT detector object 35 | sift = cv2.SIFT_create() 36 | 37 | cameras = [args.main_cam_id, 1-args.main_cam_id] 38 | for which_cam in cameras: 39 | # Read, blur, and resize the undistorted image 40 | udt_img = cv2.imread(data_root+f"{which_cam}_0_udt.png") 41 | udt_img = cv2.GaussianBlur(udt_img, (roi[4], 1), 0) 42 | udt_img = cv2.GaussianBlur(udt_img, (1, roi[5]), 0) 43 | udt_img = cv2.resize(udt_img,(udt_img.shape[1]//args.down_size,udt_img.shape[0]//args.down_size), cv2.INTER_LINEAR) 44 | cv2.imwrite(data_root+f"cam_{which_cam}.png", udt_img) 45 | 46 | frame.append(udt_img) 47 | 48 | # Detect keypoints and compute descriptors using SIFT 49 | kp, des = sift.detectAndCompute(udt_img, None) 50 | keys.append(kp) 51 | descriptors.append(des) 52 | 53 | # Draw keypoints on the image and save it 54 | tmp_frame_keys = cv2.drawKeypoints(udt_img, kp, None) 55 | 56 | cv2.imwrite(data_root+f"cam_{which_cam}_keys.png", tmp_frame_keys) 57 | 58 | ratio = 0.50 59 | matcher = cv2.BFMatcher() 60 | 61 | # Perform KNN matching of descriptors between the two images 62 | raw_matches = matcher.knnMatch(descriptors[0], descriptors[1], k = 2) 63 | 64 | # Apply ratio test to select good matches 65 | good_matches = [] 66 | for m1, m2 in raw_matches: 67 | if m1.distance < ratio * m2.distance: 68 | good_matches.append([m1]) 69 | 70 | matches = cv2.drawMatchesKnn(frame[0], keys[0], frame[1], keys[1], good_matches, None, flags = 2) 71 | cv2.imwrite(data_root+"matches.png", matches) 72 | 73 | if len(good_matches) > 4: 74 | # Extract point coordinates from the good matches 75 | ptsA = np.float32([keys[0][m[0].queryIdx].pt for m in good_matches]).reshape(-1, 1, 2) 76 | ptsB = np.float32([keys[1][m[0].trainIdx].pt for m in good_matches]).reshape(-1, 1, 2) 77 | 78 | # Compute the homography matrix using RANSAC 79 | ransacReprojThreshold = 4 80 | 81 | H, status = cv2.findHomography(ptsA, ptsB, cv2.RANSAC, ransacReprojThreshold) 82 | 83 | # Warp the second image to align with the first image 84 | img_warp = cv2.warpPerspective(frame[1], H, (frame[0].shape[1],frame[0].shape[0]),flags=cv2.INTER_LINEAR + cv2.WARP_INVERSE_MAP) 85 | cv2.imwrite(data_root+f"cam_{1-args.main_cam_id}_warp.png", img_warp) 86 | H.astype(np.float64).tofile(output_root+"H.bin") 87 | 88 | # Extract and save the cropped regions from the aligned images 89 | w_start, h_start, w_end, h_end = roi[:4] 90 | img_0 = frame[0][h_start:h_end,w_start:w_end,:] 91 | img_1 = img_warp[h_start:h_end,w_start:w_end,:] 92 | 93 | cv2.imwrite(data_root+"image0_crop.png", img_0) 94 | cv2.imwrite(data_root+"image1_crop.png", img_1) 95 | 96 | -------------------------------------------------------------------------------- /data_processing/parallel-patchmatch-master/main.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | from patchmatch import * 4 | import torch 5 | import time 6 | from sift_flow_torch import SiftFlowTorch 7 | from third_party.flowiz import flowiz 8 | 9 | import argparse 10 | import cv2 11 | import os 12 | from datetime import datetime 13 | 14 | 15 | if __name__ == "__main__": 16 | parser = argparse.ArgumentParser(description='Calibrate sample plane.') 17 | parser.add_argument('--data_root',type=str) 18 | parser.add_argument('--patch_size',type=int,default=3) 19 | parser.add_argument('--search_radius',type=int,default=50) 20 | parser.add_argument('--jump_radius',type=int,default=50) 21 | parser.add_argument('--iterations',type=int,default=10) 22 | 23 | parser.add_argument('--main_cam_id',type=int,default=0) 24 | parser.add_argument('--gpu_id',type=int,default=0) 25 | parser.add_argument('--save_root',type=str) 26 | args = parser.parse_args() 27 | 28 | save_root = args.data_root + args.save_root 29 | data_root = args.data_root + "sfm/" 30 | os.makedirs(save_root,exist_ok=True) 31 | 32 | time_start = datetime.now() 33 | patch_size = args.patch_size 34 | search_radius = args.search_radius 35 | jump_radius = args.jump_radius 36 | iterations = args.iterations 37 | 38 | img_ori = cv2.imread(data_root+'image0_crop.png') 39 | ref_ori = cv2.imread(data_root+'image1_crop.png') 40 | 41 | imgs = [img_ori, ref_ori] 42 | 43 | sift_step_size = 1 44 | sift_flow = SiftFlowTorch( 45 | cell_size=20, 46 | step_size=sift_step_size, 47 | is_boundary_included=True, 48 | num_bins=8, 49 | cuda=True, 50 | fp16=True, 51 | return_numpy=False 52 | ) 53 | device = torch.device("cuda:{}".format(args.gpu_id)) 54 | 55 | img = torch.from_numpy(img_ori).to(device) 56 | ref = torch.from_numpy(ref_ori).to(device) 57 | print('Warm-up step, will be slow on GPU') 58 | torch.cuda.synchronize() 59 | start = time.perf_counter() 60 | descs = sift_flow.extract_descriptor(imgs) 61 | torch.cuda.synchronize() 62 | end = time.perf_counter() 63 | print('Time: {:.03f} ms'.format((end - start) * 1000)) 64 | 65 | src_img = descs[0].unsqueeze(0) 66 | ref_img = descs[1].unsqueeze(0) 67 | 68 | height = src_img.shape[2] 69 | width = ref_img.shape[3] 70 | initial_NNF = torch.meshgrid(torch.arange(height),torch.arange(width)) 71 | initial_NNF = torch.stack([initial_NNF[0],initial_NNF[1]],dim=-1) 72 | initial_NNF = initial_NNF + torch.randint(-30,30,initial_NNF.shape) 73 | 74 | src_h = height - patch_size + 1 75 | src_w = width - patch_size + 1 76 | initial_NNF[:,:,0]=torch.clamp(initial_NNF[:,:,0],0, src_h-1) 77 | initial_NNF[:,:,1]=torch.clamp(initial_NNF[:,:,1],0, src_w-1) 78 | initial_NNF = initial_NNF.to(device) 79 | 80 | pm=PatchMatch(src_img,ref_img,patch_size,initial_NNF=initial_NNF,device=device) 81 | nnf = pm.run(num_iters=iterations, rand_search_radius=search_radius, jump_radius=jump_radius, allow_diagonals=True) 82 | 83 | recon_img = pm.reconstruct_without_avg(ref.permute(2,0,1).unsqueeze(0).to(src_img.dtype)) 84 | 85 | cv2.imwrite(data_root+"recon.png", recon_img) 86 | 87 | nnf = nnf.cpu().numpy().astype(np.int32) 88 | nnf.astype(np.int32).tofile(save_root+"warp.bin") 89 | 90 | 91 | 92 | 93 | 94 | -------------------------------------------------------------------------------- /data_processing/parallel-patchmatch-master/run.sh: -------------------------------------------------------------------------------- 1 | 2 | for variable in 18 3 | do 4 | class_name=paper 5 | 6 | formatted_variable=$(printf "%04d" $variable) 7 | data_root=../../database_data/"$class_name""$formatted_variable"/ 8 | 9 | texture_resolution=1024 10 | save_root=output/texture_"$texture_resolution"/ 11 | 12 | main_cam_id=0 13 | patch_size=7 14 | search_radius=50 15 | jump_radius=20 16 | iterations=10 17 | 18 | gpu_id=0 19 | 20 | python main.py --data_root $data_root --patch_size $patch_size --search_radius $search_radius --jump_radius $jump_radius --iterations $iterations --main_cam_id $main_cam_id --gpu_id $gpu_id --save_root $save_root 21 | 22 | done -------------------------------------------------------------------------------- /data_processing/parallel-patchmatch-master/third_party/flowiz/LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019 George Gach 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /data_processing/parallel-patchmatch-master/third_party/flowiz/README.md: -------------------------------------------------------------------------------- 1 | Flowiz code gotten from [https://github.com/georgegach/flowiz](https://github.com/georgegach/flowiz). -------------------------------------------------------------------------------- /data_processing/parallel-patchmatch-master/third_party/flowiz/flowiz.py: -------------------------------------------------------------------------------- 1 | # Converts Flow .flo files to Images 2 | 3 | # Author : George Gach (@georgegach) 4 | # Date : July 2019 5 | 6 | # Adapted from the Middlebury Vision project's Flow-Code 7 | # URL : http://vision.middlebury.edu/flow/ 8 | 9 | import numpy as np 10 | import os 11 | import errno 12 | from tqdm import tqdm 13 | from PIL import Image 14 | import io 15 | 16 | TAG_FLOAT = 202021.25 17 | flags = { 18 | 'debug': False 19 | } 20 | 21 | 22 | def read_flow(path): 23 | if not isinstance(path, io.BufferedReader): 24 | if not isinstance(path, str): 25 | raise AssertionError( 26 | "Input [{p}] is not a string".format(p=path)) 27 | if not os.path.isfile(path): 28 | raise AssertionError( 29 | "Path [{p}] does not exist".format(p=path)) 30 | if not path.split('.')[-1] == 'flo': 31 | raise AssertionError( 32 | "File extension [flo] required, [{f}] given".format(f=path.split('.')[-1])) 33 | 34 | flo = open(path, 'rb') 35 | else: 36 | flo = path 37 | 38 | tag = np.frombuffer(flo.read(4), np.float32, count=1)[0] 39 | if not TAG_FLOAT == tag: 40 | raise AssertionError("Wrong Tag [{t}]".format(t=tag)) 41 | 42 | width = np.frombuffer(flo.read(4), np.int32, count=1)[0] 43 | if not (width > 0 and width < 100000): 44 | raise AssertionError("Illegal width [{w}]".format(w=width)) 45 | 46 | height = np.frombuffer(flo.read(4), np.int32, count=1)[0] 47 | if not (width > 0 and width < 100000): 48 | raise AssertionError("Illegal height [{h}]".format(h=height)) 49 | 50 | nbands = 2 51 | tmp = np.frombuffer(flo.read(nbands * width * height * 4), 52 | np.float32, count=nbands * width * height) 53 | flow = np.resize(tmp, (int(height), int(width), int(nbands))) 54 | flo.close() 55 | 56 | return flow 57 | 58 | 59 | def _color_wheel(): 60 | # Original inspiration: http://members.shaw.ca/quadibloc/other/colint.htm 61 | 62 | RY = 15 63 | YG = 6 64 | GC = 4 65 | CB = 11 66 | BM = 13 67 | MR = 6 68 | 69 | ncols = RY + YG + GC + CB + BM + MR 70 | 71 | colorwheel = np.zeros([ncols, 3]) # RGB 72 | 73 | col = 0 74 | 75 | # RY 76 | colorwheel[0:RY, 0] = 255 77 | colorwheel[0:RY, 1] = np.floor(255*np.arange(0, RY, 1)/RY) 78 | col += RY 79 | 80 | # YG 81 | colorwheel[col: YG + col, 0] = 255 - \ 82 | np.floor(255*np.arange(0, YG, 1)/YG) 83 | colorwheel[col: YG + col, 1] = 255 84 | col += YG 85 | 86 | # GC 87 | colorwheel[col: GC + col, 1] = 255 88 | colorwheel[col: GC + col, 2] = np.floor(255*np.arange(0, GC, 1)/GC) 89 | col += GC 90 | 91 | # CB 92 | colorwheel[col: CB + col, 1] = 255 - \ 93 | np.floor(255*np.arange(0, CB, 1)/CB) 94 | colorwheel[col: CB + col, 2] = 255 95 | col += CB 96 | 97 | # BM 98 | colorwheel[col: BM + col, 2] = 255 99 | colorwheel[col: BM + col, 0] = np.floor(255*np.arange(0, BM, 1)/BM) 100 | col += BM 101 | 102 | # MR 103 | colorwheel[col: MR + col, 2] = 255 - \ 104 | np.floor(255*np.arange(0, MR, 1)/MR) 105 | colorwheel[col: MR + col, 0] = 255 106 | 107 | return colorwheel 108 | 109 | 110 | def _compute_color(u, v): 111 | colorwheel = _color_wheel() 112 | idxNans = np.where(np.logical_or( 113 | np.isnan(u), 114 | np.isnan(v) 115 | )) 116 | u[idxNans] = 0 117 | v[idxNans] = 0 118 | 119 | ncols = colorwheel.shape[0] 120 | radius = np.sqrt(np.multiply(u, u) + np.multiply(v, v)) 121 | a = np.arctan2(-v, -u) / np.pi 122 | fk = (a+1) / 2 * (ncols - 1) 123 | k0 = fk.astype(np.uint8) 124 | k1 = k0 + 1 125 | k1[k1 == ncols] = 0 126 | f = fk - k0 127 | 128 | img = np.empty([k1.shape[0], k1.shape[1], 3]) 129 | ncolors = colorwheel.shape[1] 130 | 131 | for i in range(ncolors): 132 | tmp = colorwheel[:, i] 133 | col0 = tmp[k0] / 255 134 | col1 = tmp[k1] / 255 135 | col = (1-f) * col0 + f * col1 136 | idx = radius <= 1 137 | col[idx] = 1 - radius[idx] * (1 - col[idx]) 138 | col[~idx] *= 0.75 139 | img[:, :, i] = np.floor(255 * col).astype(np.uint8) # RGB 140 | # img[:, :, 2 - i] = np.floor(255 * col).astype(np.uint8) # BGR 141 | 142 | return img.astype(np.uint8) 143 | 144 | 145 | def _normalize_flow(flow): 146 | UNKNOWN_FLOW_THRESH = 1e9 147 | # UNKNOWN_FLOW = 1e10 148 | 149 | height, width, nBands = flow.shape 150 | if not nBands == 2: 151 | raise AssertionError("Image must have two bands. [{h},{w},{nb}] shape given instead".format( 152 | h=height, w=width, nb=nBands)) 153 | 154 | u = flow[:, :, 0] 155 | v = flow[:, :, 1] 156 | 157 | # Fix unknown flow 158 | idxUnknown = np.where(np.logical_or( 159 | abs(u) > UNKNOWN_FLOW_THRESH, 160 | abs(v) > UNKNOWN_FLOW_THRESH 161 | )) 162 | u[idxUnknown] = 0 163 | v[idxUnknown] = 0 164 | 165 | maxu = max([-999, np.max(u)]) 166 | maxv = max([-999, np.max(v)]) 167 | minu = max([999, np.min(u)]) 168 | minv = max([999, np.min(v)]) 169 | 170 | rad = np.sqrt(np.multiply(u, u) + np.multiply(v, v)) 171 | maxrad = max([-1, np.max(rad)]) 172 | 173 | if flags['debug']: 174 | print("Max Flow : {maxrad:.4f}. Flow Range [u, v] -> [{minu:.3f}:{maxu:.3f}, {minv:.3f}:{maxv:.3f}] ".format( 175 | minu=minu, minv=minv, maxu=maxu, maxv=maxv, maxrad=maxrad 176 | )) 177 | 178 | eps = np.finfo(np.float32).eps 179 | u = u/(maxrad + eps) 180 | v = v/(maxrad + eps) 181 | 182 | return u, v 183 | 184 | 185 | def _flow2color(flow): 186 | 187 | u, v = _normalize_flow(flow) 188 | img = _compute_color(u, v) 189 | 190 | # TO-DO 191 | # Indicate unknown flows on the image 192 | # Originally done as 193 | # 194 | # IDX = repmat(idxUnknown, [1 1 3]); 195 | # img(IDX) = 0; 196 | 197 | return img 198 | 199 | 200 | def _flow2uv(flow): 201 | u, v = _normalize_flow(flow) 202 | uv = (np.dstack([u, v])*127.999+128).astype('uint8') 203 | return uv 204 | 205 | 206 | def _save_png(arr, path): 207 | # TO-DO: No dependency 208 | Image.fromarray(arr).save(path) 209 | 210 | 211 | def convert_from_file(path, mode='RGB'): 212 | return convert_from_flow(read_flow(path), mode) 213 | 214 | 215 | def convert_from_flow(flow, mode='RGB'): 216 | if mode == 'RGB': 217 | return _flow2color(flow) 218 | if mode == 'UV': 219 | return _flow2uv(flow) 220 | 221 | return _flow2color(flow) 222 | 223 | 224 | def convert_files(files, outdir=None): 225 | if outdir != None and not os.path.exists(outdir): 226 | try: 227 | os.makedirs(outdir) 228 | print("> Created directory: " + outdir) 229 | except OSError as exc: 230 | if exc.errno != errno.EEXIST: 231 | raise 232 | 233 | t = tqdm(files) 234 | for f in t: 235 | image = convert_from_file(f) 236 | 237 | if outdir == None: 238 | path = f + '.png' 239 | t.set_description(path) 240 | _save_png(image, path) 241 | else: 242 | path = os.path.join(outdir, os.path.basename(f) + '.png') 243 | t.set_description(path) 244 | _save_png(image, path) 245 | -------------------------------------------------------------------------------- /data_processing/torch_renderer/camera.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple 2 | import cv2 3 | import numpy as np 4 | import os.path as osp 5 | import torch 6 | import torch.nn.functional as F 7 | 8 | 9 | class Camera(object): 10 | 11 | def __init__( 12 | self, 13 | config_dir: str, 14 | low_res: bool = False 15 | ) -> None: 16 | """ 17 | Args: 18 | config_dir: config directory which contains cam_pos.bin, intrinsic.yml, extrinsic.yml 19 | low_res: low camera resolution. If `True`, the image size will be divided by 4. It's 20 | usually for optix simulation. 21 | """ 22 | self.cam_pos = self._load_cam_pos(osp.join(config_dir, "cam_pos.bin")) 23 | self.intrinsic, self.image_size = self._load_intrinsic(osp.join(config_dir, "intrinsic.yml"), low_res) 24 | self.extrinsic = self._load_extrinsic(osp.join(config_dir, "extrinsic.yml")) 25 | 26 | def project_points( 27 | self, 28 | position: torch.Tensor, 29 | image_coordinates: bool, 30 | ) -> torch.Tensor: 31 | """ 32 | Project the world position to camera coordinates / image 33 | coordinates. 34 | 35 | Args: 36 | position: a tensor of shape (batch, 3) 37 | image_coordinates: If `True`, unprojects the points 38 | to image coordinates. If `False` unprojects to 39 | the camera view coordinates. 40 | 41 | Returns: 42 | image_coord: a tensor of shape (batch, 2) 43 | """ 44 | device = position.device 45 | 46 | world_points = F.pad(position, (0, 1), value=1).to(device).T 47 | 48 | # world space to camera space 49 | camera_points = torch.mm(self.get_extrinsic(device), world_points) 50 | 51 | if not image_coordinates: 52 | return camera_points.T[..., :-1] 53 | 54 | # camera space to image space 55 | img_points = torch.mm(self.get_intrinsic(device), camera_points[:3, :]) 56 | 57 | img_points = img_points.transpose(0, 1) # (batch, 3) 58 | img_points = img_points[:, :2] / img_points[:, 2:3] 59 | img_points = img_points / self.get_img_size(device) 60 | 61 | return img_points 62 | 63 | def unproject_points( 64 | self, 65 | xy_depth: torch.Tensor, 66 | world_coordinates: bool 67 | ) -> torch.Tensor: 68 | """ 69 | Transform input points from camera coodinates (screen) 70 | to the world / camera coordinates. 71 | 72 | Each of the input points `xy_depth` of shape (batch, 3) is 73 | a concatenation of the x, y location and its depth. 74 | 75 | For instance, for an input 2D tensor of shape `(num_points, 3)` 76 | `xy_depth` takes the following form: 77 | `xy_depth[i] = [x[i], y[i], depth[i]]`, 78 | for a each point at an index `i`. 79 | 80 | Args: 81 | xy_depth: torch tensor of shape (batch, 3). 82 | world_coordinates: If `True`, unprojects the points back to world 83 | coordinates using the camera extrinsics `R` and `T`. 84 | `False` ignores `R` and `T` and unprojects to 85 | the camera view coordinates. 86 | 87 | Returns 88 | new_points: unprojected points with the same shape as `xy_depth`. 89 | """ 90 | device = xy_depth.device 91 | 92 | img_points = torch.cat([xy_depth[:, :2] * xy_depth[:, [2]], xy_depth[:, [2]]], dim=1) 93 | 94 | # image space to camera space 95 | camera_points = torch.mm(torch.inverse(self.get_intrinsic(device)), img_points.T) 96 | 97 | if not world_coordinates: 98 | return camera_points 99 | 100 | camera_points = F.pad(camera_points.T, (0, 1), value=1).T 101 | world_points = torch.mm(torch.inverse(self.get_extrinsic(device)), camera_points) 102 | return world_points.T[..., :-1] 103 | 104 | def _load_cam_pos(self, config_file: str) -> np.array: 105 | """ 106 | Load camera position from config file. 107 | 108 | Args: 109 | config_file: The config file is a bin file, and data of shape 110 | (3, ), dtype=float32 111 | 112 | Returns: 113 | cam_pos: shape (3, ) 114 | """ 115 | cam_pos = np.fromfile(config_file, np.float32) 116 | return cam_pos 117 | 118 | def _load_intrinsic( 119 | self, 120 | config_file: str, 121 | low_res: bool = False 122 | ) -> Tuple[np.array, np.array]: 123 | """ 124 | Load camera intrinsic 125 | 126 | Args: 127 | config_file: opencv yml file 128 | 129 | Returns: 130 | intrinsic: (3, 3) 131 | image_size: W, H 132 | """ 133 | fs2 = cv2.FileStorage(config_file, cv2.FileStorage_READ) 134 | intrinsic = fs2.getNode('camera_matrix').mat() 135 | H = fs2.getNode('image_height').real() 136 | W = fs2.getNode('image_width').real() 137 | 138 | if low_res: 139 | intrinsic = intrinsic / 4 140 | intrinsic[2, 2] = 1 141 | H = H // 4 142 | W = W // 4 143 | 144 | return intrinsic, np.array([W, H]) 145 | 146 | def _load_extrinsic(self, config_file: str) -> np.array: 147 | """ 148 | Load camera extrinsic 149 | 150 | Args: 151 | config_file: opencv yml file 152 | 153 | Returns: 154 | extrinsic: (4, 4) 155 | """ 156 | fs2 = cv2.FileStorage(config_file, cv2.FileStorage_READ) 157 | extrinsic = np.zeros((4, 4)) 158 | extrinsic[:3, :3] = fs2.getNode('rmat').mat() 159 | extrinsic[:3, 3] = fs2.getNode('tvec').mat()[:, 0] 160 | extrinsic[3, 3] = 1 161 | return extrinsic 162 | 163 | def get_intrinsic(self, device : str) -> torch.Tensor: 164 | """ 165 | Returns: 166 | intrinsic: a tensor of shape (3, 3) 167 | """ 168 | intrinsic = torch.from_numpy(self.intrinsic).float().to(device) 169 | return intrinsic 170 | 171 | def get_extrinsic(self, device : str) -> torch.Tensor: 172 | """ 173 | Returns: 174 | intrinsic: a tensor of shape (4, 4) 175 | """ 176 | extrinsic = torch.from_numpy(self.extrinsic).float().to(device) 177 | return extrinsic 178 | 179 | def get_img_size(self, device : str) -> torch.Tensor: 180 | """ 181 | Returns: 182 | image_size: a tensor of shape (2, ) 183 | """ 184 | return torch.from_numpy(self.image_size).to(device) 185 | 186 | def get_cam_pos(self, device : str) -> torch.Tensor: 187 | """ 188 | Returns: 189 | cam_pos: shape (3, ) 190 | """ 191 | return torch.from_numpy(self.cam_pos).to(device) 192 | 193 | def xy_to_raydir(self, xy_grid: torch.Tensor) -> torch.Tensor: 194 | """ 195 | Convert the `xy_grid` input of shape `(batch, 2)` to raydir. 196 | 197 | Args: 198 | xy_grid: torch.tensor grid of image xy coords. 199 | 200 | Returns: 201 | raydir: raydir of each image point in world space. 202 | torch.tensor of shape (batch, 3) 203 | """ 204 | device = xy_grid.device 205 | xy_depth = F.pad(xy_grid, (0, 1), value=1) 206 | 207 | world_pts = self.unproject_points(xy_depth, True) 208 | ray_dir = F.normalize(world_pts - self.get_cam_pos(device), dim=1) 209 | return ray_dir -------------------------------------------------------------------------------- /data_processing/torch_renderer/materials/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/OpenSVBRDF/OpenSVBRDF_source_code/27d60e6e95bc66f65d1780b114f2aa228992e7f4/data_processing/torch_renderer/materials/__init__.py -------------------------------------------------------------------------------- /data_processing/torch_renderer/materials/ggx_brdf.py: -------------------------------------------------------------------------------- 1 | import math 2 | from threading import local 3 | import torch 4 | import torch.nn.functional as F 5 | from materials.utils import cosine_sample_hemisphere 6 | from materials.utils import reflect 7 | from onb import ONB 8 | 9 | 10 | class GGX_BRDF(object): 11 | 12 | @classmethod 13 | def sample( 14 | cls, 15 | local_onb: ONB, 16 | wo: torch.Tensor, 17 | sample_num: int, 18 | ax: torch.Tensor, 19 | ay: torch.Tensor, 20 | ): 21 | """ 22 | Args: 23 | local_onb: 24 | wo: outgoing light dir in world space, shape (batch, 3) 25 | sample_num: 26 | ax: shape (batch, 1) 27 | ay: shape (batch, 1) 28 | 29 | Returns: 30 | wi: sampled direction in world space, shape (batch, N, 3) 31 | is_diff: (batch, N), if == 0, sample pd, else sample ps 32 | """ 33 | batch = wo.size(0) 34 | N = sample_num 35 | device = ax.device 36 | 37 | normal = local_onb.w() 38 | 39 | # decide to sample base on pd or ps 40 | is_diff = torch.rand(batch, N).to(device) 41 | is_diff[is_diff <= 0.5] = 0 42 | is_diff[is_diff > 0.5] = 1 43 | 44 | # sample pd 45 | unit_wi = cosine_sample_hemisphere((batch, N)).to(device) 46 | onb = ONB(batch) 47 | onb.build_from_w(normal) 48 | pd_wi = onb.inverse_transform(unit_wi) # wi in world space 49 | 50 | # sample ps 51 | z1 = torch.rand(batch, N).to(device) 52 | z2 = torch.rand(batch, N).to(device) 53 | x = ax * torch.sqrt(z1) / torch.sqrt(1 - z1) * torch.cos(2 * math.pi * z2) 54 | y = ay * torch.sqrt(z1) / torch.sqrt(1 - z1) * torch.sin(2 * math.pi * z2) 55 | z = torch.ones(batch, N).to(device) 56 | local_vhalf = F.normalize(torch.stack([-x, -y, z], dim=2), dim=2) 57 | world_half = local_onb.inverse_transform(local_vhalf) 58 | ps_wi = reflect(wo, world_half) 59 | 60 | is_diff_scale = is_diff.unsqueeze(2).repeat(1, 1, 3) 61 | wi = torch.where(is_diff_scale == 0, pd_wi, ps_wi) 62 | return wi, is_diff 63 | 64 | @classmethod 65 | def pdf( 66 | cls, 67 | local_onb: ONB, 68 | wi: torch.Tensor, 69 | wo: torch.Tensor, 70 | ax: torch.Tensor, 71 | ay: torch.Tensor, 72 | is_diff: torch.Tensor, 73 | ): 74 | """ 75 | Args: 76 | local_onb: 77 | wi: (batch, N, 3) 78 | wo: (batch, 3) 79 | ax: shape (batch, 1) 80 | ay: shape (batch, 1) 81 | is_diff: (batch, N) 82 | 83 | Returns: 84 | pdf: (batch, N) 85 | """ 86 | N = wi.size(1) 87 | 88 | # pd 89 | normalized_wi = F.normalize(wi, dim=2) 90 | local_wi = local_onb.transform(normalized_wi) 91 | pdf_pd = local_wi[:, :, 2] / math.pi 92 | 93 | # ps 94 | local_wi = F.normalize(local_onb.transform(wi), dim=2) 95 | local_wo = F.normalize(local_onb.transform(wo), dim=1) 96 | 97 | local_half = F.normalize(local_wi + local_wo.unsqueeze(1).repeat(1, N, 1), dim=2) 98 | vhalf = local_half.clone() 99 | vhalf[:, :, 0] = vhalf[:, :, 0] / ax 100 | vhalf[:, :, 1] = vhalf[:, :, 1] / ay 101 | 102 | len2 = torch.sum(vhalf * vhalf, dim=2) 103 | D = 1.0 / (math.pi * ax * ay * len2 * len2) 104 | pdf_ps = D * local_half[:, :, 2] / 4 / torch.sum(local_half * local_wi, dim=2) 105 | 106 | pdf_ps = torch.where(local_wi[:, :, 2] > 0, pdf_ps, torch.zeros_like(pdf_ps).to(pdf_ps.device)) 107 | 108 | pdf = torch.where(is_diff == 0, pdf_pd, pdf_ps) 109 | return pdf 110 | 111 | @classmethod 112 | def eval( 113 | cls, 114 | local_onb: ONB, 115 | wi: torch.Tensor, 116 | wo: torch.Tensor, 117 | ax: torch.Tensor, 118 | ay: torch.Tensor, 119 | pd: torch.Tensor, 120 | ps: torch.Tensor, 121 | is_diff: torch.Tensor = None, 122 | specular_component: str = "D_F_G_B", 123 | ) -> torch.Tensor: 124 | """ 125 | Evaluate brdf in shading coordinate(ntb space). 126 | 127 | For each shading point, we will build a coordinate system to calculate 128 | the brdf. The coordinate is usually expressed as ntb. n is the normal 129 | of the shading point. 130 | 131 | Args: 132 | local_onb: a coordinate system to do shade 133 | wi: incident light in world space, of shape (batch, lightnum, 3) 134 | wo: outgoing light in world space, of shape (batch, 3) 135 | ax: shape (batch, 1) 136 | ay: shape (batch, 1) 137 | pd: shape (batch, channel), range [0, 1] 138 | ps: shape (batch, channel), range [0, 10] 139 | is_diff: shape (batch, lightnum), if the value is 0, eval "pd_only", else if 140 | the value is 1, eval "ps_only". If `is_diff` is None, means "both" 141 | specular_component: the ingredient of BRDF, usually "D_F_G_B", B means bottom 142 | 143 | Returns: 144 | brdf: (batch, lightnum, channel) 145 | meta: 146 | """ 147 | N = wi.size(1) 148 | 149 | # transform wi and wo to local frame 150 | wi_local = local_onb.transform(wi) 151 | wo_local = local_onb.transform(wo) 152 | 153 | meta = {} 154 | 155 | a = torch.unsqueeze(pd / math.pi, dim=1) 156 | b = cls.ggx_brdf_aniso(wi_local, wo_local, ax, ay, specular_component) 157 | ps = torch.unsqueeze(ps, dim=1) 158 | 159 | if is_diff is None: 160 | brdf = a + b * ps 161 | else: 162 | is_diff_ = is_diff.unsqueeze(2) 163 | brdf = a.repeat(1, N, 1) * (1 - is_diff_) + b * ps * is_diff_ 164 | 165 | meta['pd'] = a 166 | meta['ps'] = b * ps 167 | 168 | return brdf, meta 169 | 170 | @classmethod 171 | def ggx_brdf_aniso( 172 | cls, 173 | wi: torch.Tensor, 174 | wo: torch.Tensor, 175 | ax: torch.Tensor, 176 | ay: torch.Tensor, 177 | specular_component: str 178 | ) -> torch.Tensor: 179 | """ 180 | Calculate anisotropy ggx brdf in shading coordinate. 181 | 182 | Args: 183 | wi: incident light in ntb space, of shape (batch, lightnum, 3) 184 | wo: emergent light in ntb space, of shape (batch, 3) 185 | ax: shape (batch, 1) 186 | ay: shape (batch, 1) 187 | specular_component: the ingredient of BRDF, usually "D_F_G_B" 188 | 189 | Returns: 190 | brdf: shape (batch, lightnum, 1) 191 | """ 192 | 193 | lightnum = wi.size(1) 194 | wo = torch.unsqueeze(wo, dim=1).repeat(1, lightnum, 1) 195 | 196 | wi_z = wi[:, :, [2]] 197 | wo_z = wo[:, :, [2]] 198 | denom = 4 * wi_z * wo_z 199 | vhalf = F.normalize(wi + wo, dim=2) 200 | 201 | tmp = torch.clamp(1.0 - torch.sum(wi * vhalf, dim=2, keepdim=True), 0, 1) 202 | F0 = 0.04 203 | Fresnel = F0 + (1 - F0) * tmp * tmp * tmp * tmp * tmp 204 | 205 | axayaz = torch.unsqueeze(torch.cat([ax, ay, torch.ones_like(ax)], dim=1), 206 | dim=1) 207 | vhalf = vhalf / (axayaz + 1e-6) 208 | vhalf_norm = torch.norm(vhalf, dim=2, keepdim=True) 209 | length = vhalf_norm * vhalf_norm 210 | D = 1.0 / (math.pi * torch.unsqueeze(ax, dim=1) * 211 | torch.unsqueeze(ay, dim=1) * length * length) 212 | 213 | G = cls.ggx_G1_aniso(wi, ax, ay, wi_z) * cls.ggx_G1_aniso(wo, ax, ay, wo_z) 214 | 215 | tmp = torch.ones_like(denom) 216 | if "D" in specular_component: 217 | tmp = tmp * D 218 | if "F" in specular_component: 219 | tmp = tmp * Fresnel 220 | if "G" in specular_component: 221 | tmp = tmp * G 222 | if "B" in specular_component: 223 | tmp = tmp / (denom + 1e-6) 224 | 225 | # some samples' wi_z/wo_z may less or equal than 0, should be 226 | # set to zero. Maybe this step is not necessary, because G is 227 | # already zero. 228 | tmp_zeros = torch.zeros_like(tmp) 229 | static_zero = torch.zeros(1, device=wi.device, dtype=torch.float32) 230 | res = torch.where(torch.le(wi_z, static_zero), tmp_zeros, tmp) 231 | res = torch.where(torch.le(wo_z, static_zero), tmp_zeros, res) 232 | 233 | return res 234 | 235 | @classmethod 236 | def ggx_G1_aniso( 237 | cls, 238 | v: torch.Tensor, 239 | ax: torch.Tensor, 240 | ay: torch.Tensor, 241 | vz: torch.Tensor 242 | ) -> torch.Tensor: 243 | """ 244 | If vz <= 0, return 0 245 | 246 | Args: 247 | v: shape (batch, lightnum, 3) 248 | ax: shape (batch, 1) 249 | ay: shape (batch, 1) 250 | vz: shape (batch, lightnum, 1) 251 | 252 | Returns: 253 | G1: shape (batch, lightnum, 1) 254 | """ 255 | axayaz = torch.cat([ax, ay, torch.ones_like(ax)], dim=1) 256 | vv = v * torch.unsqueeze(axayaz, dim=1) 257 | G1 = 2.0 * vz / (vz + torch.norm(vv, dim=2, keepdim=True) + 1e-6) 258 | 259 | G1 = torch.where( 260 | torch.le(vz, torch.zeros_like(vz)), 261 | torch.zeros_like(vz), 262 | G1 263 | ) 264 | return G1 -------------------------------------------------------------------------------- /data_processing/torch_renderer/materials/utils.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple 2 | import torch 3 | import torch.nn.functional as F 4 | import math 5 | 6 | 7 | def cosine_sample_hemisphere( 8 | size: Tuple 9 | ) -> torch.Tensor: 10 | 11 | u1 = torch.rand(*size, 1) 12 | u2 = torch.rand(*size, 1) 13 | 14 | # Uniformly sample disk. 15 | r = torch.sqrt(u1) 16 | phi = 2 * math.pi * u2 17 | x = r * torch.cos(phi) 18 | y = r * torch.sin(phi) 19 | 20 | # Project up to hemisphere. 21 | z = torch.sqrt(1 - x * x - y * y) 22 | 23 | return torch.cat([x, y, z], dim=len(size)) 24 | 25 | 26 | def reflect( 27 | i_: torch.Tensor, 28 | n_: torch.Tensor, 29 | ) -> torch.Tensor: 30 | 31 | assert(len(i_.size()) == 2) 32 | assert(len(n_.size()) == 3) 33 | n = F.normalize(n_, dim=2) 34 | 35 | i = i_.unsqueeze(1) 36 | dot = torch.sum(n * i, dim=2, keepdim=True) 37 | 38 | return 2 * dot * n - i 39 | -------------------------------------------------------------------------------- /data_processing/torch_renderer/onb.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | 4 | 5 | class ONB(object): 6 | 7 | def __init__(self, batch_size : int): 8 | self.batch_size = batch_size 9 | self.axis = torch.zeros(batch_size, 3, 3) 10 | self.axis[:, 0, 0] = 1 11 | self.axis[:, 1, 1] = 1 12 | self.axis[:, 2, 2] = 1 13 | 14 | def u(self) -> torch.Tensor: 15 | """ 16 | Returns: 17 | u axis of shape (batch, 3) 18 | """ 19 | return self.axis[:, 0, :] 20 | 21 | def v(self) -> torch.Tensor: 22 | return self.axis[:, 1, :] 23 | 24 | def w(self) -> torch.Tensor: 25 | return self.axis[:, 2, :] 26 | 27 | def inverse_transform(self, p : torch.Tensor) -> torch.Tensor: 28 | """ 29 | Convert local coordinate(in onb) back to global 30 | coordinate(onb in). 31 | 32 | Args: 33 | p: local coordinates of shape (batch, 3) or (batch, N, 3) 34 | 35 | Returns: 36 | global coordinates of shape (batch, 3) or (batch, N, 3) 37 | """ 38 | assert(self.batch_size == p.size(0)) 39 | assert(len(p.size()) in [2, 3]) 40 | if len(p.size()) == 2: 41 | return p[:, 0:1] * self.u() + p[:, 1:2] * self.v() + p[:, 2:3] * self.w() 42 | elif len(p.size()) == 3: 43 | u = self.u().unsqueeze(1) 44 | v = self.v().unsqueeze(1) 45 | w = self.w().unsqueeze(1) 46 | return p[:, :, [0]] * u + p[:, :, [1]] * v + p[:, :, [2]] * w 47 | 48 | def transform(self, p : torch.Tensor) -> torch.Tensor: 49 | """ 50 | Convert global coordinate(onb in) to local 51 | coordinate(in onb). 52 | 53 | Args: 54 | p: global coordinates of shape (batch, 3) or (batch, lightnum, 3) 55 | 56 | Returns: 57 | local coordinates of shape (batch, 3) or (batch, lightnum, 3) 58 | """ 59 | assert(self.batch_size == p.size(0)) 60 | assert(len(p.size()) in [2, 3]) 61 | if len(p.size()) == 2: 62 | x = torch.sum(p * self.u(), dim=1, keepdim=True) 63 | y = torch.sum(p * self.v(), dim=1, keepdim=True) 64 | z = torch.sum(p * self.w(), dim=1, keepdim=True) 65 | return torch.cat([x, y, z], dim=1) 66 | elif len(p.size()) == 3: 67 | lightnum = p.size(1) 68 | u = self.u().unsqueeze(1).repeat(1, lightnum, 1) 69 | v = self.v().unsqueeze(1).repeat(1, lightnum, 1) 70 | w = self.w().unsqueeze(1).repeat(1, lightnum, 1) 71 | x = torch.sum(p * u, dim=2, keepdim=True) 72 | y = torch.sum(p * v, dim=2, keepdim=True) 73 | z = torch.sum(p * w, dim=2, keepdim=True) 74 | return torch.cat([x, y, z], dim=2) 75 | 76 | def build_from_ntb( 77 | self, 78 | n : torch.Tensor, 79 | t : torch.Tensor, 80 | b : torch.Tensor, 81 | ) -> None: 82 | """ 83 | Args: 84 | n, t, b: The local frame, of shape (batch, 3) 85 | """ 86 | batch_size = n.size(0) 87 | self.axis = torch.zeros((batch_size, 3, 3)).to(n.device) 88 | self.axis[:, 2, :] = F.normalize(n, dim=1) 89 | self.axis[:, 1, :] = F.normalize(b, dim=1) 90 | self.axis[:, 0, :] = F.normalize(t, dim=1) 91 | 92 | def build_from_w(self, normal : torch.Tensor) -> None: 93 | """ 94 | Build the local frame based on the normal. 95 | 96 | Args: 97 | normal: The normal coordinates of shape (batch, 3) 98 | """ 99 | assert(self.batch_size == normal.size(0)) 100 | device = normal.device 101 | n = F.normalize(normal, dim=1) 102 | nz = n[:, [2]] 103 | batch_size = n.shape[0] 104 | 105 | constant_001 = torch.zeros_like(normal).to(device) 106 | constant_001[:, 2] = 1.0 107 | constant_100 = torch.zeros_like(normal).to(device) 108 | constant_100[:, 0] = 1.0 109 | 110 | nz_notequal_1 = torch.gt(torch.abs(nz - 1.0), 1e-6) 111 | nz_notequal_m1 = torch.gt(torch.abs(nz + 1.0), 1e-6) 112 | 113 | t = torch.where(nz_notequal_1 & nz_notequal_m1, constant_001, constant_100) 114 | t = F.normalize(torch.cross(t, normal), dim=1) 115 | b = torch.cross(n, t) 116 | 117 | self.axis = torch.zeros((batch_size, 3, 3)).to(device) 118 | self.axis[:, 2, :] = n 119 | self.axis[:, 1, :] = b 120 | self.axis[:, 0, :] = t 121 | 122 | def rotate_frame(self, theta : torch.Tensor) -> None: 123 | """ 124 | Rotate local frame along the normal axis 125 | 126 | Args: 127 | theta: the degrees of counterclockwise rotation, of shape (batch, 1) 128 | """ 129 | assert(self.batch_size == theta.size(0)) 130 | n = self.w() 131 | t = self.u() 132 | b = self.v() 133 | 134 | t = F.normalize(t * torch.cos(theta) + b * torch.sin(theta), dim=1) 135 | b = F.normalize(torch.cross(n, t), dim=1) 136 | self.axis = torch.zeros((self.batch_size, 3, 3)).to(theta.device) 137 | self.axis[:, 0, :] = t 138 | self.axis[:, 1, :] = b 139 | self.axis[:, 2, :] = n 140 | 141 | def _back_hemi_octa_map(self, n_2d : torch.Tensor) -> torch.Tensor: 142 | """ 143 | The original normal is (0, 0, 1), we should use this method to 144 | perturb the original normal to get a new normal and then build 145 | a new local frame based on the new normal. 146 | 147 | Args: 148 | n_2d: shape (batch, 2) 149 | 150 | Returns: 151 | local_n: shape (batch, 3), which is define in geometry local 152 | frame. 153 | """ 154 | p = (n_2d - 0.5) * 2.0 155 | resultx = (p[:, [0]] + p[:, [1]]) * 0.5 156 | resulty = (p[:, [1]] - p[:, [0]]) * 0.5 157 | resultz = 1.0 - torch.abs(resultx) - torch.abs(resulty) 158 | result = torch.cat([resultx, resulty, resultz], dim=1) 159 | return F.normalize(result, dim=1) 160 | 161 | def build_from_n2d(self, n_2d : torch.Tensor, theta : torch.Tensor) -> None: 162 | """ 163 | Args: 164 | n_2d: tensor of shape (batch, 2). the param defines how 165 | to perturb local normal. 166 | theta: the degrees of rotation of tangent. 167 | """ 168 | assert(self.batch_size == n_2d.size(0)) 169 | 170 | local_n = self._back_hemi_octa_map(n_2d) 171 | self.build_from_w(local_n) 172 | self.rotate_frame(theta) 173 | -------------------------------------------------------------------------------- /data_processing/torch_renderer/ray_trace.py: -------------------------------------------------------------------------------- 1 | import trimesh 2 | import torch 3 | import torch.nn.functional as F 4 | import numpy as np 5 | 6 | 7 | class RayTrace(object): 8 | ''' 9 | RayTrace encapsulates ray tracing operations for a given object in valid volume. 10 | 11 | functions: 12 | intersects_any: determine whether the given rays intersect with the object. 13 | 14 | intersects_location: get the position and the UV coordinates where the rays intersect with the object. 15 | 16 | ''' 17 | def __init__( 18 | self, 19 | mesh: trimesh.Trimesh, 20 | ) -> None: 21 | """ 22 | Args: 23 | mesh: an object in valid volume 24 | """ 25 | self.mesh = mesh 26 | 27 | def intersects_any( 28 | self, 29 | ray_origins: torch.Tensor, 30 | ray_dirs: torch.Tensor, 31 | ) -> torch.Tensor: 32 | """ 33 | Args: 34 | ray_origins: (batch, 3) 35 | ray_dirs: (batch, N, 3), each origin has N directions 36 | 37 | Returns: 38 | hit: (batch, N) 39 | """ 40 | batch = ray_origins.size(0) 41 | N = ray_dirs.size(1) 42 | 43 | ray_dirs = F.normalize(ray_dirs, dim=2) 44 | 45 | ray_origins_scale = torch.repeat_interleave(ray_origins, N, dim=0) 46 | 47 | ray_ori = ray_origins_scale.cpu().numpy() 48 | ray_dir = ray_dirs.reshape(-1, 3).cpu().numpy() 49 | 50 | ray_ori = ray_ori + ray_dir * 0.01 51 | 52 | hit = self.mesh.ray.intersects_any( 53 | ray_origins=ray_ori, 54 | ray_directions=ray_dir 55 | ) 56 | 57 | hit = torch.from_numpy(hit.reshape(batch, N)).to(ray_origins.device) 58 | return hit 59 | 60 | def intersects_location( 61 | self, 62 | ray_origins: torch.Tensor, 63 | ray_dirs: torch.Tensor, 64 | ) -> torch.Tensor: 65 | """ 66 | Args: 67 | ray_origins: (batch, 3) 68 | ray_dirs: (batch, N, 3), each origin has N directions 69 | Returns: 70 | hit_point: (batch, N, 3), If not hit, the position is (0, 0, 0) 71 | hit_uv: (batch, N, 2), If not hit, the uv is (0, 0) 72 | """ 73 | batch = ray_origins.size(0) 74 | N = ray_dirs.size(1) 75 | 76 | ray_dirs = F.normalize(ray_dirs, dim=2) 77 | 78 | ray_origins_scale = torch.repeat_interleave(ray_origins, N, dim=0) 79 | 80 | ray_ori = ray_origins_scale.cpu().numpy() 81 | ray_dir = F.normalize(ray_dirs, dim=2).view(-1, 3).cpu().numpy() 82 | 83 | ray_ori = ray_ori + ray_dir * 0.01 84 | 85 | locations, index_ray, index_tri = self.mesh.ray.intersects_location( 86 | ray_origins=ray_ori, 87 | ray_directions=ray_dir, 88 | multiple_hits=False, 89 | ) 90 | 91 | # Get the barycentric coordinates of the points in their respective triangle 92 | barys = trimesh.triangles.points_to_barycentric(self.mesh.vertices[self.mesh.faces[index_tri]], locations, method='cramer') 93 | uvs = np.einsum('ij,ijk->ik', barys, self.mesh.visual.uv[self.mesh.faces[index_tri]]) 94 | 95 | hit_point = np.zeros_like(ray_ori) 96 | hit_point[index_ray] = locations 97 | hit_point = torch.from_numpy(hit_point.reshape(batch, N, 3)).to(ray_origins.device).float() 98 | 99 | hit_uv = np.zeros((ray_ori.shape[0], 2)) 100 | hit_uv[index_ray] = uvs 101 | hit_uv = torch.from_numpy(hit_uv.reshape(batch, N, 2)).to(ray_origins.device).float() 102 | return hit_point, hit_uv -------------------------------------------------------------------------------- /data_processing/torch_renderer/render_utils.py: -------------------------------------------------------------------------------- 1 | from typing import Optional 2 | import torch 3 | import torch.nn.functional as F 4 | 5 | 6 | def compute_form_factors_utils(position: torch.Tensor, 7 | n: torch.Tensor, 8 | light_poses: torch.Tensor, 9 | light_normals: torch.Tensor, 10 | with_cos: Optional[bool] = True) -> torch.Tensor: 11 | """ 12 | Compute form factors in world space 13 | Formula: (wi * np) * (-wi * nl) / ||xl - xp||^2 14 | 15 | Args: 16 | position: the position of the point, of shape (batch, 3) 17 | n: the normal of the shading point, of shape (batch, 3) 18 | light_poses: shape (lightnum, 3) 19 | light_normals: shape (lightnum, 3) 20 | with_cos: if true, form factor add cos(ldir · light_normals) 21 | 22 | Returns: 23 | form_factor: (batch, lightnum, 1) 24 | """ 25 | ldir = torch.unsqueeze(light_poses, dim=0) - torch.unsqueeze(position, 26 | dim=1) 27 | dist = torch.sqrt(torch.sum(ldir**2, dim=2, 28 | keepdim=True)) 29 | ldir = F.normalize(ldir, dim=2) 30 | 31 | a = torch.sum(ldir * torch.unsqueeze(n, dim=1), dim=2, keepdim=True) 32 | a = torch.clamp(a, min=0) 33 | 34 | if not with_cos: 35 | return a 36 | 37 | b = dist * dist 38 | 39 | c = torch.sum(ldir * torch.unsqueeze(light_normals, dim=0), 40 | dim=2, 41 | keepdim=True) 42 | c = torch.clamp(c, min=0) 43 | 44 | return a / (b + 1e-6) * c 45 | -------------------------------------------------------------------------------- /data_processing/torch_renderer/setup_config.py: -------------------------------------------------------------------------------- 1 | import os 2 | import cv2 3 | from typing import Tuple, Dict 4 | import torch 5 | import numpy as np 6 | from camera import Camera 7 | 8 | 9 | class SetupConfig(object): 10 | """ 11 | This class manages the setup of lightstage. 12 | 13 | The setup should be saved in `wallet_of_torch_renderer` folder, include many 14 | configs: `cam_pos`, `lights_data`, `lights_angular`, `lights_intensity`, 15 | `visualize_map`, `rgb_tensor` and so on. 16 | """ 17 | 18 | def __init__( 19 | self, 20 | config_dir: str, 21 | mask: torch.Tensor = None, 22 | low_res: bool = False, 23 | ) -> None: 24 | """ 25 | Args: 26 | config_dir: the directory of the config 27 | mask: a bool tensor to remove unwanted lights. The shape is 28 | (lightnum, ), if true, the light will be kept, else 29 | the light will be removed. 30 | low_res: low camera resolution, usually for optix simulation. 31 | """ 32 | self.config_dir = config_dir 33 | self.mask = mask 34 | 35 | self.camera = Camera(self.config_dir, low_res) 36 | 37 | self.light_poses, self.light_normals = self._load_light_data( 38 | self.config_dir + "lights.bin") 39 | 40 | self.lights_angular = self._load_lights_angular( 41 | self.config_dir + "lights_angular.npy") 42 | 43 | self.lights_intensity = self._load_lights_intensity( 44 | self.config_dir + "lights_intensity.npy") 45 | 46 | self.img_size, self.visualize_map = self._load_visualize_map( 47 | self.config_dir + "visualize_config_torch.bin") 48 | 49 | self.rgb_tensor = self._load_rgb_tensor( 50 | self.config_dir + "color_tensor.bin" 51 | ) 52 | 53 | self.mask_data = self._load_mask_data( 54 | self.config_dir + "mask.npy" 55 | ) 56 | 57 | self.trained_idx = self._load_trained_idx( 58 | self.config_dir + "trained_idx.txt" 59 | ) 60 | 61 | def _load_light_data( 62 | self, config_file: str) -> Tuple[torch.Tensor, torch.Tensor]: 63 | """ 64 | Load light position and normal from config file. 65 | 66 | TODO: Attention! The normal of lights is reversed! You should generate new setup config. 67 | 68 | Args: 69 | config_file: The config file is a bin file, and has data of shape 70 | (2, lightnum, 3), dtype=float32. In the first dim, 0 is position 71 | and 1 is normal. 72 | 73 | Returns: 74 | light_poses: a ndarray of shape (lightnum, 3) 75 | light_normals: a ndarray of shape (lightnum, 3) 76 | """ 77 | light_data = np.fromfile(config_file, np.float32).reshape([2, -1, 3]) 78 | light_poses, light_normals = light_data[0], light_data[1] 79 | if self.mask is not None: 80 | light_poses = light_data[0, self.mask, :] 81 | light_normals = light_data[1, self.mask, :] 82 | return light_poses, light_normals 83 | 84 | def _load_lights_angular(self, config_file: str) -> torch.Tensor: 85 | """ 86 | Load angular distribution parameters of light intensity. 87 | 88 | The distribution is fitted with two cubic polynomials, therefore each 89 | light is with 8 parameters. See `compute_light_distribution` function 90 | for more detail. 91 | 92 | Args: 93 | config_file: a npy file. 94 | 95 | Returns: 96 | lights_angular: a ndarray of shape (1, lightnum, 8) 97 | """ 98 | lights_angular = None 99 | if os.path.isfile(config_file): 100 | lights_angular = np.load(config_file) 101 | if self.mask is not None: 102 | lights_angular = lights_angular[0, self.mask, :] 103 | return lights_angular 104 | 105 | def _load_lights_intensity(self, config_file: str) -> torch.Tensor: 106 | """ 107 | Load the relative intensity of lights. 108 | 109 | Args: 110 | config_file: a npy file 111 | 112 | Returns: 113 | lights_intensity: a ndarray of shape (lightnum, 3). 114 | """ 115 | lights_intensity = None 116 | if os.path.isfile(config_file): 117 | lights_intensity = np.load(config_file) 118 | if self.mask is not None: 119 | lights_intensity = lights_intensity[self.mask, :] 120 | return lights_intensity 121 | 122 | def _load_visualize_map( 123 | self, config_file: str) -> Tuple[torch.Tensor, torch.Tensor]: 124 | """ 125 | Load visualize mapping data from config file. 126 | 127 | The sequence of lights in `light_poses` data. for example: our 128 | light board has 48*64 lights. The lights sequence in `light_poses` 129 | starts from right bottom corner, from bottom to top, from left to 130 | right. Then the map idx sequence will be 63 47 63 46 63 45 ... 131 | 63 0 62 47 62 46 ... 0 0. 132 | 133 | Args: 134 | config_file: The config file is a bin file with dtype int32. The 135 | first 2 integers are image size (W, H), after are map idx. 136 | 137 | Returns: 138 | img_size: a ndarray of shape (2, ), represents H, W 139 | visualize_map: a ndarray of shape (lightnum, 2) 140 | 141 | """ 142 | with open(config_file, "rb") as pf: 143 | img_size = np.fromfile(pf, np.int32, 2) 144 | visualize_map = np.fromfile(pf, np.int32).reshape([-1, 2]) 145 | return img_size, visualize_map 146 | 147 | def _load_rgb_tensor(self, config_file): 148 | """ 149 | Load rgb tensor from config file 150 | 151 | Args: 152 | config_file: 153 | 154 | Returns: 155 | rgb_tensor: a ndarray of shape (3, 3, 3), (light, object, camera) 156 | """ 157 | rgb_tensor = np.zeros((3, 3, 3)) 158 | for i in range(3): 159 | rgb_tensor[i][i][i] = 1.0 160 | 161 | if os.path.isfile(config_file): 162 | rgb_tensor = np.fromfile(config_file, np.float32).reshape([3, 3, 3]) 163 | else: 164 | print("Note: color tensor file is not found, use default color tensor!") 165 | 166 | return rgb_tensor 167 | 168 | def _load_mask_data(self, config_file): 169 | """ 170 | Load mask data from config file 171 | 172 | Args: 173 | config_file: a npy file 174 | 175 | Returns: 176 | A dict contians `mask_anchor`, `mask_axis_x`, `mask_axis_y` 177 | """ 178 | mask_data = np.load(config_file, allow_pickle=True).item() 179 | return mask_data 180 | 181 | def _load_trained_idx(self, config_file): 182 | if os.path.exists(config_file): 183 | trained_idx = np.loadtxt(config_file) 184 | else: 185 | trained_idx = [] 186 | return trained_idx 187 | 188 | def print_configs(self) -> None: 189 | """ 190 | Print the configs 191 | """ 192 | print("[SETUP CONFIG]") 193 | print("cam_pos:", self.cam_pos) 194 | 195 | def get_rgb_tensor(self, device : str) -> torch.Tensor: 196 | """ 197 | Returns: 198 | rgb_tensor: a tensor of shape (3, 3, 3), (light, object, camera) 199 | """ 200 | return torch.from_numpy(self.rgb_tensor).to(device) 201 | 202 | def get_cam_pos(self, device : str) -> torch.Tensor: 203 | """ 204 | Returns: 205 | cam_pos: shape (3, ) 206 | """ 207 | return self.camera.get_cam_pos(device) 208 | 209 | def get_lights_intensity(self, device : str) -> torch.Tensor: 210 | """ 211 | Returns: 212 | lights_intensity: a tensor of shape (lightnum, 3). 213 | """ 214 | return torch.from_numpy(self.lights_intensity).to(device) 215 | 216 | def get_lights_angular(self, device : str) -> torch.Tensor: 217 | """ 218 | Returns: 219 | lights_angular: a tensor of shape (1, lightnum, 8) 220 | """ 221 | return torch.from_numpy(self.lights_angular).to(device) 222 | 223 | def get_light_normals(self, device : str) -> torch.Tensor: 224 | """ 225 | Returns: 226 | light_normals: a tensor of shape (lightnum, 3) 227 | """ 228 | return torch.from_numpy(self.light_normals).to(device) 229 | 230 | def get_light_poses(self, device : str) -> torch.Tensor: 231 | """ 232 | Returns: 233 | light_poses: a tensor of shape (lightnum, 3) 234 | """ 235 | return torch.from_numpy(self.light_poses).to(device) 236 | 237 | def get_light_idx(self, light_pos: np.array) -> int: 238 | """ 239 | Get the nearest light 240 | 241 | Args: 242 | light_pos: ndarray, (1, 3) 243 | """ 244 | assert(light_pos.shape[0] == 1 and light_pos.shape[1] == 3) 245 | dist = self.light_poses - light_pos 246 | light_idx = np.argmin(np.sum(dist * dist, axis=1)) 247 | return light_idx 248 | 249 | def get_trained_idx(self) -> list: 250 | """ 251 | Returns: 252 | trained_idx: a list of trained_idx 253 | """ 254 | return self.trained_idx 255 | 256 | def get_vis_img_size(self) -> Tuple[int, int]: 257 | """ 258 | Returns: 259 | H, W 260 | """ 261 | return self.img_size[0], self.img_size[1] 262 | 263 | def get_visualize_map(self, device : str) -> torch.Tensor: 264 | """ 265 | Returns: 266 | visualize_map: a tensor of shape (lightnum, 2) 267 | """ 268 | return torch.from_numpy(self.visualize_map).long().to(device) 269 | 270 | def get_mask_data(self, device : str) -> Dict: 271 | """ 272 | Returns: 273 | mask_anchor: 274 | mask_axis_x: 275 | mask_axis_y: 276 | """ 277 | anchor = torch.from_numpy(self.mask_data['mask_anchor']).float().to(device) 278 | offset1 = torch.from_numpy(self.mask_data['mask_axis_x']).float().to(device) 279 | offset2 = torch.from_numpy(self.mask_data['mask_axis_y']).float().to(device) 280 | return anchor, offset1, offset2 281 | 282 | def get_light_num(self) -> int: 283 | return self.light_normals.shape[0] 284 | 285 | @classmethod 286 | def convert_xy_to_idx(cls, x, y): 287 | """ 288 | Args: 289 | x, y: left top is (0, 0), x is vertical 290 | 291 | Returns: 292 | idx: light index, right buttom is 0, right top is 47. 293 | """ 294 | return 48 * (63 - y) + (47 - x) 295 | 296 | @classmethod 297 | def convert_idx_to_xy(cls, idx): 298 | x = 47 - idx % 48 299 | y = 63 - idx // 48 300 | return x, y 301 | 302 | @classmethod 303 | def get_downsampled_light_poses(cls, ori_light_poses): 304 | """ 305 | Args: 306 | light_poses: (3072, 3) 307 | Returns: 308 | light_poses downsampled 309 | """ 310 | light_poses = np.zeros((3072 // 16, 3)) 311 | cnt = 0 312 | for row in range(48 // 4): 313 | for col in range(64 // 4): 314 | ori_idx = cls.convert_xy_to_idx(4*row+2, 4*col+2) 315 | light_poses[cnt] = ori_light_poses[ori_idx] 316 | cnt += 1 317 | return light_poses 318 | 319 | @classmethod 320 | def upsample_data( 321 | cls, 322 | light_data: np.array 323 | ): 324 | """ 325 | Args: 326 | light_data: (192, *) 327 | """ 328 | new_light_data = np.zeros((3072, light_data.shape[1])) 329 | for i in range(light_data.shape[0]): 330 | row = i // 16 331 | col = i % 16 332 | for x in range(4*row, 4*(row+1)): 333 | for y in range(4*col, 4*(col+1)): 334 | idx = cls.convert_xy_to_idx(x, y) 335 | new_light_data[idx] = light_data[i] 336 | return new_light_data -------------------------------------------------------------------------------- /data_processing/torch_renderer/wallet_of_torch_renderer/lightstage/cam_pos.bin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/OpenSVBRDF/OpenSVBRDF_source_code/27d60e6e95bc66f65d1780b114f2aa228992e7f4/data_processing/torch_renderer/wallet_of_torch_renderer/lightstage/cam_pos.bin -------------------------------------------------------------------------------- /data_processing/torch_renderer/wallet_of_torch_renderer/lightstage/color_tensor.bin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/OpenSVBRDF/OpenSVBRDF_source_code/27d60e6e95bc66f65d1780b114f2aa228992e7f4/data_processing/torch_renderer/wallet_of_torch_renderer/lightstage/color_tensor.bin -------------------------------------------------------------------------------- /data_processing/torch_renderer/wallet_of_torch_renderer/lightstage/extrinsic.yml: -------------------------------------------------------------------------------- 1 | %YAML:1.0 2 | --- 3 | calibration_time: "Fri May 20 18:56:25 2022" 4 | rmat: !!opencv-matrix 5 | rows: 3 6 | cols: 3 7 | dt: f 8 | data: [ -0.0015425 , 0.99992293, -0.0123192, 9 | 0.99999164, 0.00158901, 0.00376663, 10 | 0.00378592, -0.01231328, -0.99991702 ] 11 | tvec: !!opencv-matrix 12 | rows: 3 13 | cols: 1 14 | dt: f 15 | data: [ -1.0899953065510286e+00, -4.7868011322118376e+00, 16 | 2.8901836360109729e+02 ] 17 | -------------------------------------------------------------------------------- /data_processing/torch_renderer/wallet_of_torch_renderer/lightstage/lights.bin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/OpenSVBRDF/OpenSVBRDF_source_code/27d60e6e95bc66f65d1780b114f2aa228992e7f4/data_processing/torch_renderer/wallet_of_torch_renderer/lightstage/lights.bin -------------------------------------------------------------------------------- /data_processing/torch_renderer/wallet_of_torch_renderer/lightstage/mask.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/OpenSVBRDF/OpenSVBRDF_source_code/27d60e6e95bc66f65d1780b114f2aa228992e7f4/data_processing/torch_renderer/wallet_of_torch_renderer/lightstage/mask.npy -------------------------------------------------------------------------------- /data_processing/torch_renderer/wallet_of_torch_renderer/lightstage/visualize_config_torch.bin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/OpenSVBRDF/OpenSVBRDF_source_code/27d60e6e95bc66f65d1780b114f2aa228992e7f4/data_processing/torch_renderer/wallet_of_torch_renderer/lightstage/visualize_config_torch.bin -------------------------------------------------------------------------------- /data_processing/visualize_results/rendering.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This is the experimental code for paper ``Ma et al., OpenSVBRDF: A Database of Measured Spatially-Varying Reflectance, SIGGRAPH Asia 2023``. 3 | This script is suboptimal and experimental. 4 | There may be redundant lines and functionalities. 5 | 6 | Xiaohe Ma, 2024/02 7 | ''' 8 | 9 | 10 | ''' 11 | This script reads in line measurements file and interpolates the pixels by lerp(). Then, it crops the result image if needed. 12 | 13 | by Leyao 14 | ''' 15 | import numpy as np 16 | import argparse 17 | import torch 18 | import sys 19 | import math 20 | import os 21 | os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1" 22 | import cv2 23 | import csv 24 | import matplotlib.pyplot as plt 25 | import matplotlib 26 | sys.path.append("../finetune/") 27 | import AUTO_planar_scanner_net_inference 28 | 29 | TORCH_RENDER_PATH = "../torch_renderer/" 30 | sys.path.append(TORCH_RENDER_PATH) 31 | # import torch_render 32 | from torch_render import TorchRender 33 | from setup_config import SetupConfig 34 | 35 | from skimage.metrics import structural_similarity as ssim 36 | import skimage 37 | 38 | RENDER_SCALAR = 3e3/math.pi 39 | 40 | 41 | if __name__ == "__main__": 42 | parser = argparse.ArgumentParser() 43 | parser.add_argument("data_root",default="") 44 | parser.add_argument("save_root",default="") 45 | parser.add_argument("model_file",type=str,default="") 46 | parser.add_argument("pattern_file",type=str,default="") 47 | parser.add_argument("--lighting_pattern_num",type=int,default=64) 48 | parser.add_argument("--tex_resolution",type=int,default=1024) 49 | parser.add_argument("--main_cam_id",type=int,default=0) 50 | parser.add_argument("--config_dir",type=str,default="../device_configuration/") 51 | parser.add_argument("--gpu_id",type=int,default=0) 52 | parser.add_argument("--batch_size",type=int,default=100) 53 | parser.add_argument("--shape_latent_len",type=int,default=48) 54 | parser.add_argument("--color_latent_len",type=int,default=8) 55 | args = parser.parse_args() 56 | 57 | root = args.data_root 58 | os.makedirs(args.save_root, exist_ok=True) 59 | 60 | compute_device = torch.device(f"cuda:{args.gpu_id}") 61 | setup = SetupConfig(TORCH_RENDER_PATH+"wallet_of_torch_renderer/lightstage/") 62 | torch_render = TorchRender(setup) 63 | 64 | 65 | m_len_perview = 3 66 | scalar = 255*0.04 67 | uvs = np.fromfile(args.data_root+"texture_1024/texturemap_uv.bin", np.int32).reshape([-1,2]) 68 | 69 | uv_map = np.fromfile(args.data_root+f"texture_{args.tex_resolution}/uvs_cam{args.main_cam_id}.bin", np.float32).astype(np.int32) 70 | 71 | uv_map = uv_map.reshape([-1,2]) 72 | uv_map = uv_map // 2 73 | min_uv = np.min(uv_map, axis=0) 74 | uv_map = uv_map - min_uv 75 | 76 | w, h = np.max(uv_map, axis=0) + 1 77 | 78 | 79 | ############# GT 80 | gt_result = np.fromfile(args.data_root+f"texture_{args.tex_resolution}/line_measurements_{args.tex_resolution}.bin",np.float32).reshape((-1,args.lighting_pattern_num,m_len_perview)) 81 | 82 | 83 | gt_imgs = np.zeros([h, w, args.lighting_pattern_num, 3], dtype=np.float32) 84 | gt_imgs[uv_map[:, 1], uv_map[:, 0]] = gt_result 85 | gt_imgs = np.transpose(gt_imgs,(2,0,1,3)) * scalar 86 | 87 | 88 | ############# LATENT 89 | train_configs = {} 90 | train_configs["training_device"] = 0 91 | train_configs["lumitexel_length"] = 64*64*3 92 | train_configs["shape_latent_len"] = args.shape_latent_len 93 | train_configs["color_latent_len"] = args.color_latent_len 94 | 95 | pretrained_dict = torch.load(args.model_file, map_location=compute_device) 96 | 97 | inference_net = AUTO_planar_scanner_net_inference.PlanarScannerNet(train_configs) 98 | print("loading trained model...") 99 | something_not_found = False 100 | model_dict = inference_net.state_dict() 101 | for k,_ in model_dict.items(): 102 | if k not in pretrained_dict: 103 | print("not found:", k) 104 | something_not_found = True 105 | if something_not_found: 106 | exit() 107 | 108 | model_dict = inference_net.state_dict() 109 | pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} 110 | model_dict.update(pretrained_dict) 111 | inference_net.load_state_dict(model_dict) 112 | for p in inference_net.parameters(): 113 | p.requires_grad=False 114 | inference_net.to(compute_device) 115 | inference_net.eval() 116 | 117 | 118 | lighting_patterns = np.fromfile(args.pattern_file,np.float32).reshape((args.lighting_pattern_num,-1,3)).astype(np.float32) 119 | lighting_patterns = lighting_patterns.reshape([1,args.lighting_pattern_num,-1,3]) 120 | lighting_patterns = torch.from_numpy(lighting_patterns).to(compute_device) 121 | 122 | latent_len = args.color_latent_len + args.shape_latent_len 123 | 124 | batch_size = args.batch_size 125 | lumitexel_length = 64*64*3 126 | all_latent_len = args.shape_latent_len + 3 * args.color_latent_len 127 | 128 | pf_latent = open(args.data_root+f"latent/pass3_latent_{args.tex_resolution}.bin","rb") 129 | 130 | pf_latent.seek(0,2) 131 | texel_num = pf_latent.tell()//all_latent_len//4 132 | pf_latent.seek(0,0) 133 | print("texel num : ", texel_num) 134 | 135 | texel_sequence = np.arange(texel_num) 136 | 137 | ptr = 0 138 | 139 | latent_result = np.zeros([texel_num,args.lighting_pattern_num,3],np.float32) 140 | latent_result = torch.from_numpy(latent_result).to(compute_device) 141 | 142 | while True: 143 | 144 | tmp_sequence = texel_sequence[ptr:ptr+batch_size] 145 | if tmp_sequence.shape[0] == 0: 146 | break 147 | tmp_seq_size = tmp_sequence.shape[0] 148 | 149 | tmp_latent = np.fromfile(pf_latent,np.float32,count=all_latent_len*tmp_seq_size).reshape([tmp_seq_size,all_latent_len]) 150 | tmp_latent = torch.from_numpy(tmp_latent).to(compute_device) 151 | 152 | color_latent = tmp_latent[:,:3 * args.color_latent_len].reshape([tmp_seq_size,3,args.color_latent_len]) 153 | shape_latent = tmp_latent[:,3 * args.color_latent_len:].reshape([tmp_seq_size,1,args.shape_latent_len]).repeat(1,3,1) 154 | 155 | color_shape_latent = torch.cat([color_latent,shape_latent],dim=-1) 156 | color_shape_latent = color_shape_latent.reshape([tmp_seq_size*3,latent_len]) 157 | 158 | _,tmp_nn_lumi = inference_net(color_shape_latent,input_is_latent=True) 159 | tmp_nn_lumi = torch.max(torch.zeros_like(tmp_nn_lumi),tmp_nn_lumi) 160 | 161 | tmp_nn_lumi = tmp_nn_lumi.reshape([tmp_seq_size,3,-1]).permute(0,2,1) 162 | 163 | 164 | tmp_measurements = torch.sum(lighting_patterns*tmp_nn_lumi.unsqueeze(dim=1),dim=2).reshape([tmp_seq_size,-1,3]) 165 | latent_result[ptr:ptr+batch_size,:,:] = tmp_measurements 166 | 167 | ptr += batch_size 168 | 169 | 170 | latent_imgs = np.zeros([h, w, args.lighting_pattern_num, 3], dtype=np.float32) 171 | latent_imgs[uv_map[:, 1], uv_map[:, 0]] = latent_result.cpu().numpy() 172 | 173 | latent_imgs = np.transpose(latent_imgs,(2,0,1,3)) * scalar 174 | 175 | ############# TEXTURE_MAPS 176 | texure_root = args.data_root+"texture_maps/" 177 | positions = cv2.imread(texure_root + "pos_texture.exr", 6)[:,:,::-1].reshape([-1,3]) 178 | fitted_axay = cv2.imread(texure_root+"ax_ay_texture.exr",6)[:,:,::-1][:,:,:2].reshape([-1,2]) 179 | fitted_normal = cv2.imread(texure_root+"normal_texture.exr",6)[:,:,::-1].reshape([-1,3]) 180 | fitted_pd = cv2.imread(texure_root+"pd_texture.exr",6)[:,:,::-1].reshape([-1,3]) 181 | fitted_ps = cv2.imread(texure_root+"ps_texture.exr",6)[:,:,::-1].reshape([-1,3]) 182 | fitted_tangent = cv2.imread(texure_root+"tangent_texture.exr",6)[:,:,::-1].reshape([-1,3]) 183 | 184 | fitted_normal = (fitted_normal - 0.5) * 2.0 185 | fitted_tangent = (fitted_tangent - 0.5) * 2.0 186 | 187 | fitted_params_rgb = np.concatenate([np.zeros([texel_num,3],np.float32), fitted_axay, fitted_pd, fitted_ps],axis=-1) 188 | 189 | texel_sequence = np.arange(texel_num) 190 | ptr = 0 191 | batch_size = args.batch_size 192 | 193 | tex_result = np.zeros([texel_num,args.lighting_pattern_num,3],np.float32) 194 | tex_result = torch.from_numpy(tex_result).to(compute_device) 195 | 196 | while True: 197 | 198 | tmp_sequence = texel_sequence[ptr:ptr+batch_size] 199 | if tmp_sequence.shape[0] == 0: 200 | break 201 | tmp_seq_size = tmp_sequence.shape[0] 202 | pos = torch.from_numpy(positions[tmp_sequence]).to(compute_device).to(torch.float32) 203 | 204 | params = fitted_params_rgb[tmp_sequence] 205 | 206 | params = torch.from_numpy(params).to(compute_device).to(torch.float32) 207 | 208 | n = torch.from_numpy(fitted_normal[tmp_sequence]).to(compute_device) 209 | t = torch.from_numpy(fitted_tangent[tmp_sequence]).to(compute_device) 210 | b = torch.cross(n,t) 211 | 212 | rotate_theta = torch.zeros(tmp_seq_size,1,device=compute_device,dtype=torch.float32) 213 | shading_frame = [n,t,b] 214 | 215 | lumi, end_points = torch_render.generate_lumitexel( 216 | params, 217 | pos, 218 | global_custom_frame=[n,t,b], 219 | use_custom_frame="ntb", 220 | pd_ps_wanted="both", 221 | ) 222 | 223 | lumi = lumi.reshape(tmp_seq_size,setup.get_light_num(),3)*RENDER_SCALAR 224 | 225 | 226 | measurements = torch.sum(lighting_patterns*lumi.unsqueeze(dim=1),dim=2).reshape([tmp_seq_size,-1,3]) 227 | 228 | tex_result[ptr:ptr+batch_size,:,:] = measurements 229 | ptr += batch_size 230 | 231 | 232 | tex_imgs = np.zeros([h, w, args.lighting_pattern_num, 3], dtype=np.float32) 233 | tex_imgs[uv_map[:, 1], uv_map[:, 0]] = tex_result.cpu().numpy() 234 | 235 | 236 | tex_imgs = np.transpose(tex_imgs,(2,0,1,3)) * scalar 237 | 238 | 239 | img_num = args.lighting_pattern_num 240 | loss = np.zeros([img_num, 2], np.float32) 241 | imgs = [] 242 | 243 | for which_img in range(img_num): 244 | gt_img = np.array(gt_imgs[which_img], np.float32) 245 | latent_img = np.array(latent_imgs[which_img], np.float32) 246 | tex_img = np.array(tex_imgs[which_img], np.float32) 247 | 248 | loss[which_img, 0] = ssim(gt_img, latent_img, data_range=255, channel_axis=-1) 249 | loss[which_img, 1] = ssim(gt_img, tex_img, data_range=255, channel_axis=-1) 250 | 251 | latent_error = np.abs(gt_img-latent_img).reshape([-1,3]) 252 | tex_error = np.abs(gt_img-tex_img).reshape([-1,3]) 253 | 254 | h,w = gt_img.shape[:2] 255 | 256 | gt_error = matplotlib.cm.jet(np.zeros_like(latent_error)/255)[:,0,:3]*255/2 257 | gt_error = gt_error.reshape([h,w,3]) 258 | 259 | latent_error = matplotlib.cm.jet(latent_error/255)[:,0,:3]*255/2 260 | latent_error = latent_error.reshape([h,w,3]) 261 | 262 | tex_error = matplotlib.cm.jet(tex_error/255)[:,0,:3]*255/2 263 | tex_error = tex_error.reshape([h,w,3]) 264 | 265 | latent_img = cv2.putText(latent_img, '%.2f' % loss[which_img, 0], (0,50),cv2.FONT_HERSHEY_SIMPLEX,2,(0,255,0),5) 266 | tex_img = cv2.putText(tex_img, '%.2f' % loss[which_img, 1], (0,50),cv2.FONT_HERSHEY_SIMPLEX,2,(0,255,0),5) 267 | 268 | photo_img = np.concatenate([gt_img, latent_img, tex_img],axis=1) 269 | error_img = np.concatenate([gt_error,latent_error,tex_error],axis=1) 270 | img = np.concatenate([photo_img, error_img],axis=0) 271 | img = img[:,:,::-1] 272 | 273 | cv2.imwrite(args.save_root+f"{which_img}.png",img) 274 | 275 | 276 | np.savetxt(args.save_root + "ssim_loss.csv", loss, delimiter=',', fmt='%.2f') 277 | 278 | 279 | 280 | 281 | 282 | 283 | 284 | 285 | 286 | 287 | 288 | 289 | -------------------------------------------------------------------------------- /data_processing/visualize_results/run.sh: -------------------------------------------------------------------------------- 1 | 2 | for variable in 18 3 | do 4 | class_name=paper 5 | formatted_variable=$(printf "%04d" $variable) 6 | data_root=../../database_data/$class_name$formatted_variable/output/ 7 | # data_root=../../../$class_name$variable/output/ 8 | save_root=$data_root/render/ 9 | 10 | model_root=../../database_model 11 | pattern_file=$model_root/opt_W_64.bin 12 | model_file=$model_root/latent_48_24_500000_2437925.pkl 13 | tex_resolution=1024 14 | 15 | gpu_id=1 16 | main_cam_id=0 17 | shape_latent_len=48 18 | color_latent_len=8 19 | lighting_pattern_num=64 20 | 21 | 22 | python rendering.py $data_root $save_root $model_file $pattern_file --tex_resolution $tex_resolution --main_cam_id $main_cam_id --lighting_pattern_num $lighting_pattern_num 23 | 24 | done 25 | -------------------------------------------------------------------------------- /database_data/download.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/OpenSVBRDF/OpenSVBRDF_source_code/27d60e6e95bc66f65d1780b114f2aa228992e7f4/database_data/download.md -------------------------------------------------------------------------------- /database_model/download.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/OpenSVBRDF/OpenSVBRDF_source_code/27d60e6e95bc66f65d1780b114f2aa228992e7f4/database_model/download.md -------------------------------------------------------------------------------- /download.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # download paper0018 4 | wget -P database_data/ https://drive.google.com/file/d/1PBhkDUvGb9goTIzc_9HxCAtOB72KGw9K/view?usp=drive_link 5 | 6 | # download satin0002 7 | wget -P database_data/ https://drive.google.com/file/d/1tKWYCvvX073X8HIgc2kjC8_QoGUM5QVP/view?usp=drive_link 8 | 9 | # download ceramic0014 10 | wget -P database_data/ https://drive.google.com/file/d/1mCItNXJprGko8PGxUIXF5GHMozbrE-rC/view?usp=drive_link 11 | 12 | # download metal0005 13 | wget -P database_data/ https://drive.google.com/file/d/1uEQoDfjriP45v-uAYOL2JdKe2IA0wHSs/view?usp=drive_link 14 | 15 | # download model 16 | wget -P . https://drive.google.com/file/d/1px3Ij1B7GIESWhAAm0-MHOhwR6yjWaVB/view?usp=drive_link 17 | 18 | # download device configuration 19 | wget -P . https://drive.google.com/file/d/1dIqEQcImBUaTGfsy0SVb8S6Pjua5u317/view?usp=drive_link 20 | 21 | unzip database_data/paper0018.zip -d database_data/ 22 | unzip database_data/satin0002.zip -d database_data/ 23 | unzip database_data/ceramic0014.zip -d database_data/ 24 | unzip database_data/metal0005.zip -d database_data/ 25 | 26 | unzip database_model.zip -d database_model/ 27 | unzip device_configuration.zip -d data_processing/device_configuration/ 28 | -------------------------------------------------------------------------------- /environment.yml: -------------------------------------------------------------------------------- 1 | name: opensvbrdf 2 | channels: 3 | - defaults 4 | dependencies: 5 | - _libgcc_mutex=0.1=main 6 | - _openmp_mutex=5.1=1_gnu 7 | - ca-certificates=2024.9.24=h06a4308_0 8 | - ld_impl_linux-64=2.40=h12ee557_0 9 | - libffi=3.4.4=h6a678d5_1 10 | - libgcc-ng=11.2.0=h1234567_1 11 | - libgomp=11.2.0=h1234567_1 12 | - libstdcxx-ng=11.2.0=h1234567_1 13 | - ncurses=6.4=h6a678d5_0 14 | - openssl=3.0.15=h5eee18b_0 15 | - pip=24.2=py38h06a4308_0 16 | - python=3.8.20=he870216_0 17 | - readline=8.2=h5eee18b_0 18 | - setuptools=75.1.0=py38h06a4308_0 19 | - sqlite=3.45.3=h5eee18b_0 20 | - tk=8.6.14=h39e8969_0 21 | - wheel=0.44.0=py38h06a4308_0 22 | - xz=5.4.6=h5eee18b_1 23 | - zlib=1.2.13=h5eee18b_1 24 | - pip: 25 | - contourpy==1.1.1 26 | - cycler==0.12.1 27 | - fonttools==4.54.1 28 | - imageio==2.35.1 29 | - imath==0.0.2 30 | - importlib-resources==6.4.5 31 | - joblib==1.4.2 32 | - kiwisolver==1.4.7 33 | - lazy-loader==0.4 34 | - matplotlib==3.7.1 35 | - networkx==3.1 36 | - numpy==1.23.5 37 | - nvidia-cublas-cu11==11.10.3.66 38 | - nvidia-cuda-nvrtc-cu11==11.7.99 39 | - nvidia-cuda-runtime-cu11==11.7.99 40 | - nvidia-cudnn-cu11==8.5.0.96 41 | - opencv-contrib-python==4.9.0.80 42 | - opencv-python==4.7.0.72 43 | - openexr==3.2.4 44 | - packaging==24.1 45 | - pillow==10.4.0 46 | - pyparsing==3.1.4 47 | - python-dateutil==2.9.0.post0 48 | - pywavelets==1.4.1 49 | - scikit-image==0.21.0 50 | - scikit-learn==1.2.2 51 | - scipy==1.10.1 52 | - six==1.16.0 53 | - threadpoolctl==3.5.0 54 | - tifffile==2023.7.10 55 | - torch==1.13.0 56 | - tqdm==4.65.0 57 | - trimesh==3.21.5 58 | - typing-extensions==4.12.2 59 | - zipp==3.20.2 60 | -------------------------------------------------------------------------------- /readme.md: -------------------------------------------------------------------------------- 1 | # OpenSVBRDF: A Database of Measured Spatially-Varying Reflectance 2 | 3 | This is the source code for our paper accepted at Siggraph Asia 2023. 4 | 5 | ![Teaser Image](assets/sorted_teaser_new.jpg) 6 | 7 | The database is available at [OpenSVBRDF](https://opensvbrdf.github.io/). Currently, all texture maps are available for download. Due to the large data volume, we are still looking for storage solutions for the neural representations and raw capture images. Currently, we provide the raw images of two samples for users to run this code and reproduce the results of the paper. 8 | 9 | ## Download Data & Model & Device Configurations 10 | 11 | We offer several samples to test the results comprehensively. Each sample is about 16 GB. You can comment out some of the lines for downloading samples in the `download.sh` file based on your available storage space. 12 | 13 | ## Running Steps 14 | 15 | We provide two approaches to run the code, you can choose any one you prefer. 16 | 17 | ### Running with Conda 18 | 19 | We tested our code on Ubuntu 22.04.3 LTS with NVIDIA Driver Version 535.129.03, and CUDA 12.2. 20 | 21 | 1. Create a new environment and install the packages. 22 | 23 | ``` 24 | conda env create -f environment.yml 25 | conda activate opensvbrdf 26 | ``` 27 | 2. Before running each command, confirm the sample name you want to process. The output of each step will be saved in the `output/` folder. 28 | 29 | #### Step-by-Step Instructions 30 | 31 | First, `cd data_processing/` to set your work directory as data_processing/. 32 | 33 | 1. `cd generate_uv/` and run `run.sh` to automatically determine the final texture map area from masks and captured photos. 34 | 2. `cd parallel-patchmatch-master/` and run `run.sh` to align photos from two cameras using dense patchmatch, finding corresponding pixel positions in the secondary camera for each pixel in the primary camera. 35 | 3. `cd extract_measurements/` and run `run.sh` to extract all measurement values for each pixel in the texture map from each photo. 36 | 4. `cd finetune/` and run `run.sh` to infer the neural representations and perform the three finetune steps described in the paper. (change the `server_num` in line 21 as the number of gpus you have) 37 | 5. `cd fitting/` and run `run.sh` to fit the neural representation into 6D PBR texture maps ready for rendering. 38 | 6. `cd visualize_results` and run `run.sh` to render the neural representations and PBR texture maps, compute SSIM, and visualize MSE error against real photos for result validation. 39 | 40 | After completing these steps, you can view the reconstructed results in the `output/render/` folder within the sample directory. 41 | 42 | ### Running with Docker 43 | 44 | 1. Similar to the previous appoach, confirm the sample name you want to process before running each command. 45 | 46 | 2. Create docker image at the root of the project 47 | 48 | `docker build -t opensvbrdf .` 49 | 50 | 3. Create docker container and enter it 51 | 52 | `docker run -it --rm --gpus all -v ./database_data/:/app/database_data opensvbrdf` 53 | 54 | 4. Run the code, and the output will be saved in the `database_data/sample_name/output/` folder 55 | 56 | `bash /app/run.sh` 57 | 58 | 59 | ## License 60 | 61 | This code is licensed under GPL-3.0. If you use our data or code in your work, please cite our paper: 62 | 63 | ``` 64 | @article{ma2023opensvbrdf, 65 | title={OpenSVBRDF: A Database of Measured Spatially-Varying Reflectance}, 66 | author={Ma, Xiaohe and Xu, Xianmin and Zhang, Leyao and Zhou, Kun and Wu, Hongzhi}, 67 | journal={ACM Transactions on Graphics (TOG)}, 68 | volume={42}, 69 | number={6}, 70 | pages={1--14}, 71 | year={2023}, 72 | publisher={ACM New York, NY, USA} 73 | } 74 | ``` 75 | 76 | We also acknowledge the use of code from [SIFT-Flow-GPU](https://github.com/hmorimitsu/sift-flow-gpu) for camera alignment. -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | Imath==0.0.2 2 | matplotlib==3.7.1 3 | numpy==1.23.5 4 | opencv_contrib_python==4.9.0.80 5 | opencv_python==4.7.0.72 6 | # OpenEXR==1.3.9 7 | OpenEXR==3.2.4 8 | # Pillow==9.4.0 9 | Pillow==10.4.0 10 | scikit_learn==1.2.2 11 | scikit-image==0.21.0 12 | torch==1.13.0 13 | tqdm==4.65.0 14 | trimesh==3.21.5 15 | -------------------------------------------------------------------------------- /run.sh: -------------------------------------------------------------------------------- 1 | cd /app/data_processing/generate_uv/ 2 | bash run.sh 3 | 4 | cd /app/data_processing/parallel-patchmatch-master/ 5 | bash run.sh 6 | 7 | cd /app/data_processing/extract_measurements/ 8 | bash run.sh 9 | 10 | cd /app/data_processing/finetune/ 11 | bash run.sh 12 | 13 | cd /app/data_processing/fitting/ 14 | bash run.sh 15 | 16 | cd /app/data_processing/visualize_results/ 17 | bash run.sh --------------------------------------------------------------------------------