Spaces:
Build error
Build error
| import torch | |
| import torch.nn as nn | |
| from torch import optim | |
| from pathlib import Path | |
| from easydict import EasyDict | |
| import pytorch_lightning as pl | |
| import torch.nn.functional as F | |
| import torchvision.transforms as T | |
| from torchvision.utils import save_image | |
| from torchmetrics import MeanSquaredError, StructuralSimilarityIndexMeasure | |
| #from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger | |
| from . import DenseReg, RenderingLoss | |
| from ..render import Renderer, encode_as_unit_interval, gamma_decode, gamma_encode | |
| class Vanilla(pl.LightningModule): | |
| metrics = ['I_mse','N_mse','A_mse','R_mse','I_ssim','N_ssim','A_ssim','R_ssim'] | |
| maps = {'I': 'reco', 'N': 'normals', 'R': 'roughness', 'A': 'albedo'} | |
| def __init__(self, model: nn.Module, loss: DenseReg = None, lr: float = 0, batch_size: int = 0, max_images: int = 10): | |
| super().__init__() | |
| self.model = model | |
| self.loss = loss | |
| self.lr = lr | |
| self.batch_size = batch_size | |
| self.tanh = nn.Tanh() | |
| self.norm = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) | |
| self.val_images = [] | |
| self.max_images = max_images | |
| self.save_hyperparameters(ignore=['model', 'loss']) | |
| def training_step(self, x): | |
| y = self(*x) | |
| loss = self.loss(x, y) | |
| self.log_to('train', loss) | |
| return dict(loss=loss.total, y=y) | |
| def forward(self, src, tgt): | |
| src_out, tgt_out = None, None | |
| if None not in src: | |
| src_out = self.model(self.norm(src.input)) | |
| self.post_process_(src_out) | |
| if None not in tgt: | |
| tgt_out = self.model(self.norm(tgt.input)) | |
| self.post_process_(tgt_out) | |
| return src_out, tgt_out | |
| def post_process_(self, o: EasyDict): | |
| # (1) activation function, (2) concat unit z, (3) normalize to unit vector | |
| nxy = self.tanh(o.normals) | |
| nx, ny = torch.split(nxy*3, split_size_or_sections=1, dim=1) | |
| n = torch.cat([nx, ny, torch.ones_like(nx)], dim=1) | |
| o.normals = F.normalize(n, dim=1) | |
| # (1) activation function, (2) mapping [-1,1]->[0,1] | |
| a = self.tanh(o.albedo) | |
| o.albedo = encode_as_unit_interval(a) | |
| # (1) activation function, (2) mapping [-1,1]->[0,1], (3) channel repeat x3 | |
| r = self.tanh(o.roughness) | |
| o.roughness = encode_as_unit_interval(r.repeat(1,3,1,1)) | |
| def validation_step(self, x, *_): | |
| y = self(*x) | |
| loss = self.loss(x, y) | |
| self.log_to('val', loss) | |
| if len(self.val_images) * self.batch_size < self.max_images and self.logger: | |
| self.val_images.append((x, y)) | |
| return dict(loss=loss.total, y=y) | |
| def on_validation_epoch_end(self): | |
| cur_ind = 0 | |
| for ind, (x, y) in enumerate(self.val_images): | |
| for key in ['normals', 'albedo', 'roughness']: | |
| for i in range(len(x[0][key])): | |
| pred_image = y[0][key][i] | |
| gt_image = x[0][key][i] | |
| self.logger.experiment.add_image(f'val/{cur_ind}_{key}_pred', pred_image, self.global_step) | |
| self.logger.experiment.add_image(f'val/{cur_ind}_{key}_gt', gt_image, self.global_step) | |
| cur_ind += 1 | |
| if cur_ind >= self.max_images: | |
| break | |
| self.val_images.clear() | |
| def log_to(self, split, loss): | |
| self.log_dict({f'{split}/{k}': v for k, v in loss.items()}, batch_size=self.batch_size) | |
| def log_images(self, x, y, split, max_images=5): | |
| # Log predicted and ground truth images | |
| for key in ['normals', 'albedo', 'roughness']: | |
| for i in range(min(len(x[0][key]), max_images)): | |
| self.logger.experiment.add_image(f'{split}/{key}_pred_{i}', y[0][key][i], self.global_step) | |
| self.logger.experiment.add_image(f'{split}/{key}_gt_{i}', x[0][key][i], self.global_step) | |
| def on_test_start(self): | |
| self.renderer = RenderingLoss(Renderer()) | |
| for m in Vanilla.metrics: | |
| if 'mse' in m: | |
| setattr(self, m, MeanSquaredError().to(self.device)) | |
| elif 'ssim' in m: | |
| setattr(self, m, StructuralSimilarityIndexMeasure(data_range=1).to(self.device)) | |
| def test_step(self, x, batch_idx, dl_id=0): | |
| y = self.model(self.norm(x.input)) | |
| self.post_process_(y) | |
| # image reconstruction | |
| y.reco = self.renderer.reconstruction(y, x.input_params) | |
| x.reco = gamma_decode(x.input) | |
| for m in Vanilla.metrics: | |
| mapid, *_ = m | |
| k = Vanilla.maps[mapid] | |
| meter = getattr(self, m) | |
| meter(y[k], x[k].to(y[k].dtype)) | |
| self.log(m, getattr(self, m), on_epoch=True) | |
| if self.logger: | |
| self.log_images(x, y, split='test') | |
| def predict_step(self, x, batch_idx): | |
| y = self.model(self.norm(x.input)) | |
| self.post_process_(y) | |
| I, name, outdir = x.input[0], x.name[0], Path(x.path[0]).parent | |
| N_pred, A_pred, R_pred = y.normals[0], y.albedo[0], y.roughness[0] | |
| save_image(gamma_encode(A_pred), outdir/f'{name}_albedo.png') | |
| save_image(encode_as_unit_interval(N_pred), outdir/f'{name}_normals.png') | |
| save_image(R_pred, outdir/f'{name}_roughness.png') | |
| def configure_optimizers(self): | |
| optimizer = optim.Adam(self.model.parameters(), lr=self.lr, weight_decay=1e-4) | |
| return dict(optimizer=optimizer) | |