Spaces:
Runtime error
Runtime error
| # MIT License | |
| # Copyright (c) 2022 Intelligent Systems Lab Org | |
| # Permission is hereby granted, free of charge, to any person obtaining a copy | |
| # of this software and associated documentation files (the "Software"), to deal | |
| # in the Software without restriction, including without limitation the rights | |
| # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell | |
| # copies of the Software, and to permit persons to whom the Software is | |
| # furnished to do so, subject to the following conditions: | |
| # The above copyright notice and this permission notice shall be included in all | |
| # copies or substantial portions of the Software. | |
| # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR | |
| # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, | |
| # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE | |
| # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER | |
| # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, | |
| # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE | |
| # SOFTWARE. | |
| # File author: Shariq Farooq Bhat | |
| import torch | |
| import torch.cuda.amp as amp | |
| import torch.nn as nn | |
| from zoedepth.trainers.loss import GradL1Loss, SILogLoss | |
| from zoedepth.utils.config import DATASETS_CONFIG | |
| from zoedepth.utils.misc import compute_metrics | |
| from zoedepth.data.preprocess import get_black_border | |
| from .base_trainer import BaseTrainer | |
| from torchvision import transforms | |
| from PIL import Image | |
| import numpy as np | |
| class Trainer(BaseTrainer): | |
| def __init__(self, config, model, train_loader, test_loader=None, device=None): | |
| super().__init__(config, model, train_loader, | |
| test_loader=test_loader, device=device) | |
| self.device = device | |
| self.silog_loss = SILogLoss() | |
| self.grad_loss = GradL1Loss() | |
| self.scaler = amp.GradScaler(enabled=self.config.use_amp) | |
| def train_on_batch(self, batch, train_step): | |
| """ | |
| Expects a batch of images and depth as input | |
| batch["image"].shape : batch_size, c, h, w | |
| batch["depth"].shape : batch_size, 1, h, w | |
| """ | |
| images, depths_gt = batch['image'].to( | |
| self.device), batch['depth'].to(self.device) | |
| dataset = batch['dataset'][0] | |
| b, c, h, w = images.size() | |
| mask = batch["mask"].to(self.device).to(torch.bool) | |
| losses = {} | |
| with amp.autocast(enabled=self.config.use_amp): | |
| output = self.model(images) | |
| pred_depths = output['metric_depth'] | |
| l_si, pred = self.silog_loss( | |
| pred_depths, depths_gt, mask=mask, interpolate=True, return_interpolated=True) | |
| loss = self.config.w_si * l_si | |
| losses[self.silog_loss.name] = l_si | |
| if self.config.w_grad > 0: | |
| l_grad = self.grad_loss(pred, depths_gt, mask=mask) | |
| loss = loss + self.config.w_grad * l_grad | |
| losses[self.grad_loss.name] = l_grad | |
| else: | |
| l_grad = torch.Tensor([0]) | |
| self.scaler.scale(loss).backward() | |
| if self.config.clip_grad > 0: | |
| self.scaler.unscale_(self.optimizer) | |
| nn.utils.clip_grad_norm_( | |
| self.model.parameters(), self.config.clip_grad) | |
| self.scaler.step(self.optimizer) | |
| if self.should_log and (self.step % int(self.config.log_images_every * self.iters_per_epoch)) == 0: | |
| # -99 is treated as invalid depth in the log_images function and is colored grey. | |
| depths_gt[torch.logical_not(mask)] = -99 | |
| self.log_images(rgb={"Input": images[0, ...]}, depth={"GT": depths_gt[0], "PredictedMono": pred[0]}, prefix="Train", | |
| min_depth=DATASETS_CONFIG[dataset]['min_depth'], max_depth=DATASETS_CONFIG[dataset]['max_depth']) | |
| if self.config.get("log_rel", False): | |
| self.log_images( | |
| scalar_field={"RelPred": output["relative_depth"][0]}, prefix="TrainRel") | |
| self.scaler.update() | |
| self.optimizer.zero_grad() | |
| return losses | |
| def eval_infer(self, x): | |
| with amp.autocast(enabled=self.config.use_amp): | |
| m = self.model.module if self.config.multigpu else self.model | |
| pred_depths = m(x)['metric_depth'] | |
| return pred_depths | |
| def crop_aware_infer(self, x): | |
| # if we are not avoiding the black border, we can just use the normal inference | |
| if not self.config.get("avoid_boundary", False): | |
| return self.eval_infer(x) | |
| # otherwise, we need to crop the image to avoid the black border | |
| # For now, this may be a bit slow due to converting to numpy and back | |
| # We assume no normalization is done on the input image | |
| # get the black border | |
| assert x.shape[0] == 1, "Only batch size 1 is supported for now" | |
| x_pil = transforms.ToPILImage()(x[0].cpu()) | |
| x_np = np.array(x_pil, dtype=np.uint8) | |
| black_border_params = get_black_border(x_np) | |
| top, bottom, left, right = black_border_params.top, black_border_params.bottom, black_border_params.left, black_border_params.right | |
| x_np_cropped = x_np[top:bottom, left:right, :] | |
| x_cropped = transforms.ToTensor()(Image.fromarray(x_np_cropped)) | |
| # run inference on the cropped image | |
| pred_depths_cropped = self.eval_infer(x_cropped.unsqueeze(0).to(self.device)) | |
| # resize the prediction to x_np_cropped's size | |
| pred_depths_cropped = nn.functional.interpolate( | |
| pred_depths_cropped, size=(x_np_cropped.shape[0], x_np_cropped.shape[1]), mode="bilinear", align_corners=False) | |
| # pad the prediction back to the original size | |
| pred_depths = torch.zeros((1, 1, x_np.shape[0], x_np.shape[1]), device=pred_depths_cropped.device, dtype=pred_depths_cropped.dtype) | |
| pred_depths[:, :, top:bottom, left:right] = pred_depths_cropped | |
| return pred_depths | |
| def validate_on_batch(self, batch, val_step): | |
| images = batch['image'].to(self.device) | |
| depths_gt = batch['depth'].to(self.device) | |
| dataset = batch['dataset'][0] | |
| mask = batch["mask"].to(self.device) | |
| if 'has_valid_depth' in batch: | |
| if not batch['has_valid_depth']: | |
| return None, None | |
| depths_gt = depths_gt.squeeze().unsqueeze(0).unsqueeze(0) | |
| mask = mask.squeeze().unsqueeze(0).unsqueeze(0) | |
| if dataset == 'nyu': | |
| pred_depths = self.crop_aware_infer(images) | |
| else: | |
| pred_depths = self.eval_infer(images) | |
| pred_depths = pred_depths.squeeze().unsqueeze(0).unsqueeze(0) | |
| with amp.autocast(enabled=self.config.use_amp): | |
| l_depth = self.silog_loss( | |
| pred_depths, depths_gt, mask=mask.to(torch.bool), interpolate=True) | |
| metrics = compute_metrics(depths_gt, pred_depths, **self.config) | |
| losses = {f"{self.silog_loss.name}": l_depth.item()} | |
| if val_step == 1 and self.should_log: | |
| depths_gt[torch.logical_not(mask)] = -99 | |
| self.log_images(rgb={"Input": images[0]}, depth={"GT": depths_gt[0], "PredictedMono": pred_depths[0]}, prefix="Test", | |
| min_depth=DATASETS_CONFIG[dataset]['min_depth'], max_depth=DATASETS_CONFIG[dataset]['max_depth']) | |
| return metrics, losses | |