Initial media depth project backup
This commit is contained in:
Codex
@@ -0,0 +1,326 @@
|
||||
# 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 os
|
||||
import uuid
|
||||
import warnings
|
||||
from datetime import datetime as dt
|
||||
from typing import Dict
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
import wandb
|
||||
from tqdm import tqdm
|
||||
|
||||
from zoedepth.utils.config import flatten
|
||||
from zoedepth.utils.misc import RunningAverageDict, colorize, colors
|
||||
|
||||
|
||||
def is_rank_zero(args):
|
||||
return args.rank == 0
|
||||
|
||||
|
||||
class BaseTrainer:
|
||||
def __init__(self, config, model, train_loader, test_loader=None, device=None):
|
||||
""" Base Trainer class for training a model."""
|
||||
|
||||
self.config = config
|
||||
self.metric_criterion = "abs_rel"
|
||||
if device is None:
|
||||
device = torch.device(
|
||||
'cuda') if torch.cuda.is_available() else torch.device('cpu')
|
||||
self.device = device
|
||||
self.model = model
|
||||
self.train_loader = train_loader
|
||||
self.test_loader = test_loader
|
||||
self.optimizer = self.init_optimizer()
|
||||
self.scheduler = self.init_scheduler()
|
||||
|
||||
def resize_to_target(self, prediction, target):
|
||||
if prediction.shape[2:] != target.shape[-2:]:
|
||||
prediction = nn.functional.interpolate(
|
||||
prediction, size=target.shape[-2:], mode="bilinear", align_corners=True
|
||||
)
|
||||
return prediction
|
||||
|
||||
def load_ckpt(self, checkpoint_dir="./checkpoints", ckpt_type="best"):
|
||||
import glob
|
||||
import os
|
||||
|
||||
from zoedepth.models.model_io import load_wts
|
||||
|
||||
if hasattr(self.config, "checkpoint"):
|
||||
checkpoint = self.config.checkpoint
|
||||
elif hasattr(self.config, "ckpt_pattern"):
|
||||
pattern = self.config.ckpt_pattern
|
||||
matches = glob.glob(os.path.join(
|
||||
checkpoint_dir, f"*{pattern}*{ckpt_type}*"))
|
||||
if not (len(matches) > 0):
|
||||
raise ValueError(f"No matches found for the pattern {pattern}")
|
||||
checkpoint = matches[0]
|
||||
else:
|
||||
return
|
||||
model = load_wts(self.model, checkpoint)
|
||||
# TODO : Resuming training is not properly supported in this repo. Implement loading / saving of optimizer and scheduler to support it.
|
||||
print("Loaded weights from {0}".format(checkpoint))
|
||||
warnings.warn(
|
||||
"Resuming training is not properly supported in this repo. Implement loading / saving of optimizer and scheduler to support it.")
|
||||
self.model = model
|
||||
|
||||
def init_optimizer(self):
|
||||
m = self.model.module if self.config.multigpu else self.model
|
||||
|
||||
if self.config.same_lr:
|
||||
print("Using same LR")
|
||||
if hasattr(m, 'core'):
|
||||
m.core.unfreeze()
|
||||
params = self.model.parameters()
|
||||
else:
|
||||
print("Using diff LR")
|
||||
if not hasattr(m, 'get_lr_params'):
|
||||
raise NotImplementedError(
|
||||
f"Model {m.__class__.__name__} does not implement get_lr_params. Please implement it or use the same LR for all parameters.")
|
||||
|
||||
params = m.get_lr_params(self.config.lr)
|
||||
|
||||
return optim.AdamW(params, lr=self.config.lr, weight_decay=self.config.wd)
|
||||
|
||||
def init_scheduler(self):
|
||||
lrs = [l['lr'] for l in self.optimizer.param_groups]
|
||||
return optim.lr_scheduler.OneCycleLR(self.optimizer, lrs, epochs=self.config.epochs, steps_per_epoch=len(self.train_loader),
|
||||
cycle_momentum=self.config.cycle_momentum,
|
||||
base_momentum=0.85, max_momentum=0.95, div_factor=self.config.div_factor, final_div_factor=self.config.final_div_factor, pct_start=self.config.pct_start, three_phase=self.config.three_phase)
|
||||
|
||||
def train_on_batch(self, batch, train_step):
|
||||
raise NotImplementedError
|
||||
|
||||
def validate_on_batch(self, batch, val_step):
|
||||
raise NotImplementedError
|
||||
|
||||
def raise_if_nan(self, losses):
|
||||
for key, value in losses.items():
|
||||
if torch.isnan(value):
|
||||
raise ValueError(f"{key} is NaN, Stopping training")
|
||||
|
||||
@property
|
||||
def iters_per_epoch(self):
|
||||
return len(self.train_loader)
|
||||
|
||||
@property
|
||||
def total_iters(self):
|
||||
return self.config.epochs * self.iters_per_epoch
|
||||
|
||||
def should_early_stop(self):
|
||||
if self.config.get('early_stop', False) and self.step > self.config.early_stop:
|
||||
return True
|
||||
|
||||
def train(self):
|
||||
print(f"Training {self.config.name}")
|
||||
if self.config.uid is None:
|
||||
self.config.uid = str(uuid.uuid4()).split('-')[-1]
|
||||
run_id = f"{dt.now().strftime('%d-%h_%H-%M')}-{self.config.uid}"
|
||||
self.config.run_id = run_id
|
||||
self.config.experiment_id = f"{self.config.name}{self.config.version_name}_{run_id}"
|
||||
self.should_write = ((not self.config.distributed)
|
||||
or self.config.rank == 0)
|
||||
self.should_log = self.should_write # and logging
|
||||
if self.should_log:
|
||||
tags = self.config.tags.split(
|
||||
',') if self.config.tags != '' else None
|
||||
wandb.init(project=self.config.project, name=self.config.experiment_id, config=flatten(self.config), dir=self.config.root,
|
||||
tags=tags, notes=self.config.notes, settings=wandb.Settings(start_method="fork"))
|
||||
|
||||
self.model.train()
|
||||
self.step = 0
|
||||
best_loss = np.inf
|
||||
validate_every = int(self.config.validate_every * self.iters_per_epoch)
|
||||
|
||||
|
||||
if self.config.prefetch:
|
||||
|
||||
for i, batch in tqdm(enumerate(self.train_loader), desc=f"Prefetching...",
|
||||
total=self.iters_per_epoch) if is_rank_zero(self.config) else enumerate(self.train_loader):
|
||||
pass
|
||||
|
||||
losses = {}
|
||||
def stringify_losses(L): return "; ".join(map(
|
||||
lambda kv: f"{colors.fg.purple}{kv[0]}{colors.reset}: {round(kv[1].item(),3):.4e}", L.items()))
|
||||
for epoch in range(self.config.epochs):
|
||||
if self.should_early_stop():
|
||||
break
|
||||
|
||||
self.epoch = epoch
|
||||
################################# Train loop ##########################################################
|
||||
if self.should_log:
|
||||
wandb.log({"Epoch": epoch}, step=self.step)
|
||||
pbar = tqdm(enumerate(self.train_loader), desc=f"Epoch: {epoch + 1}/{self.config.epochs}. Loop: Train",
|
||||
total=self.iters_per_epoch) if is_rank_zero(self.config) else enumerate(self.train_loader)
|
||||
for i, batch in pbar:
|
||||
if self.should_early_stop():
|
||||
print("Early stopping")
|
||||
break
|
||||
# print(f"Batch {self.step+1} on rank {self.config.rank}")
|
||||
losses = self.train_on_batch(batch, i)
|
||||
# print(f"trained batch {self.step+1} on rank {self.config.rank}")
|
||||
|
||||
self.raise_if_nan(losses)
|
||||
if is_rank_zero(self.config) and self.config.print_losses:
|
||||
pbar.set_description(
|
||||
f"Epoch: {epoch + 1}/{self.config.epochs}. Loop: Train. Losses: {stringify_losses(losses)}")
|
||||
self.scheduler.step()
|
||||
|
||||
if self.should_log and self.step % 50 == 0:
|
||||
wandb.log({f"Train/{name}": loss.item()
|
||||
for name, loss in losses.items()}, step=self.step)
|
||||
|
||||
self.step += 1
|
||||
|
||||
########################################################################################################
|
||||
|
||||
if self.test_loader:
|
||||
if (self.step % validate_every) == 0:
|
||||
self.model.eval()
|
||||
if self.should_write:
|
||||
self.save_checkpoint(
|
||||
f"{self.config.experiment_id}_latest.pt")
|
||||
|
||||
################################# Validation loop ##################################################
|
||||
# validate on the entire validation set in every process but save only from rank 0, I know, inefficient, but avoids divergence of processes
|
||||
metrics, test_losses = self.validate()
|
||||
# print("Validated: {}".format(metrics))
|
||||
if self.should_log:
|
||||
wandb.log(
|
||||
{f"Test/{name}": tloss for name, tloss in test_losses.items()}, step=self.step)
|
||||
|
||||
wandb.log({f"Metrics/{k}": v for k,
|
||||
v in metrics.items()}, step=self.step)
|
||||
|
||||
if (metrics[self.metric_criterion] < best_loss) and self.should_write:
|
||||
self.save_checkpoint(
|
||||
f"{self.config.experiment_id}_best.pt")
|
||||
best_loss = metrics[self.metric_criterion]
|
||||
|
||||
self.model.train()
|
||||
|
||||
if self.config.distributed:
|
||||
dist.barrier()
|
||||
# print(f"Validated: {metrics} on device {self.config.rank}")
|
||||
|
||||
# print(f"Finished step {self.step} on device {self.config.rank}")
|
||||
#################################################################################################
|
||||
|
||||
# Save / validate at the end
|
||||
self.step += 1 # log as final point
|
||||
self.model.eval()
|
||||
self.save_checkpoint(f"{self.config.experiment_id}_latest.pt")
|
||||
if self.test_loader:
|
||||
|
||||
################################# Validation loop ##################################################
|
||||
metrics, test_losses = self.validate()
|
||||
# print("Validated: {}".format(metrics))
|
||||
if self.should_log:
|
||||
wandb.log({f"Test/{name}": tloss for name,
|
||||
tloss in test_losses.items()}, step=self.step)
|
||||
wandb.log({f"Metrics/{k}": v for k,
|
||||
v in metrics.items()}, step=self.step)
|
||||
|
||||
if (metrics[self.metric_criterion] < best_loss) and self.should_write:
|
||||
self.save_checkpoint(
|
||||
f"{self.config.experiment_id}_best.pt")
|
||||
best_loss = metrics[self.metric_criterion]
|
||||
|
||||
self.model.train()
|
||||
|
||||
def validate(self):
|
||||
with torch.no_grad():
|
||||
losses_avg = RunningAverageDict()
|
||||
metrics_avg = RunningAverageDict()
|
||||
for i, batch in tqdm(enumerate(self.test_loader), desc=f"Epoch: {self.epoch + 1}/{self.config.epochs}. Loop: Validation", total=len(self.test_loader), disable=not is_rank_zero(self.config)):
|
||||
metrics, losses = self.validate_on_batch(batch, val_step=i)
|
||||
|
||||
if losses:
|
||||
losses_avg.update(losses)
|
||||
if metrics:
|
||||
metrics_avg.update(metrics)
|
||||
|
||||
return metrics_avg.get_value(), losses_avg.get_value()
|
||||
|
||||
def save_checkpoint(self, filename):
|
||||
if not self.should_write:
|
||||
return
|
||||
root = self.config.save_dir
|
||||
if not os.path.isdir(root):
|
||||
os.makedirs(root)
|
||||
|
||||
fpath = os.path.join(root, filename)
|
||||
m = self.model.module if self.config.multigpu else self.model
|
||||
torch.save(
|
||||
{
|
||||
"model": m.state_dict(),
|
||||
"optimizer": None, # TODO : Change to self.optimizer.state_dict() if resume support is needed, currently None to reduce file size
|
||||
"epoch": self.epoch
|
||||
}, fpath)
|
||||
|
||||
def log_images(self, rgb: Dict[str, list] = {}, depth: Dict[str, list] = {}, scalar_field: Dict[str, list] = {}, prefix="", scalar_cmap="jet", min_depth=None, max_depth=None):
|
||||
if not self.should_log:
|
||||
return
|
||||
|
||||
if min_depth is None:
|
||||
try:
|
||||
min_depth = self.config.min_depth
|
||||
max_depth = self.config.max_depth
|
||||
except AttributeError:
|
||||
min_depth = None
|
||||
max_depth = None
|
||||
|
||||
depth = {k: colorize(v, vmin=min_depth, vmax=max_depth)
|
||||
for k, v in depth.items()}
|
||||
scalar_field = {k: colorize(
|
||||
v, vmin=None, vmax=None, cmap=scalar_cmap) for k, v in scalar_field.items()}
|
||||
images = {**rgb, **depth, **scalar_field}
|
||||
wimages = {
|
||||
prefix+"Predictions": [wandb.Image(v, caption=k) for k, v in images.items()]}
|
||||
wandb.log(wimages, step=self.step)
|
||||
|
||||
def log_line_plot(self, data):
|
||||
if not self.should_log:
|
||||
return
|
||||
|
||||
plt.plot(data)
|
||||
plt.ylabel("Scale factors")
|
||||
wandb.log({"Scale factors": wandb.Image(plt)}, step=self.step)
|
||||
plt.close()
|
||||
|
||||
def log_bar_plot(self, title, labels, values):
|
||||
if not self.should_log:
|
||||
return
|
||||
|
||||
data = [[label, val] for (label, val) in zip(labels, values)]
|
||||
table = wandb.Table(data=data, columns=["label", "value"])
|
||||
wandb.log({title: wandb.plot.bar(table, "label",
|
||||
"value", title=title)}, step=self.step)
|
||||
@@ -0,0 +1,48 @@
|
||||
# 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
|
||||
|
||||
from importlib import import_module
|
||||
|
||||
|
||||
def get_trainer(config):
|
||||
"""Builds and returns a trainer based on the config.
|
||||
|
||||
Args:
|
||||
config (dict): the config dict (typically constructed using utils.config.get_config)
|
||||
config.trainer (str): the name of the trainer to use. The module named "{config.trainer}_trainer" must exist in trainers root module
|
||||
|
||||
Raises:
|
||||
ValueError: If the specified trainer does not exist under trainers/ folder
|
||||
|
||||
Returns:
|
||||
Trainer (inherited from zoedepth.trainers.BaseTrainer): The Trainer object
|
||||
"""
|
||||
assert "trainer" in config and config.trainer is not None and config.trainer != '', "Trainer not specified. Config: {0}".format(
|
||||
config)
|
||||
try:
|
||||
Trainer = getattr(import_module(
|
||||
f"zoedepth.trainers.{config.trainer}_trainer"), 'Trainer')
|
||||
except ModuleNotFoundError as e:
|
||||
raise ValueError(f"Trainer {config.trainer}_trainer not found.") from e
|
||||
return Trainer
|
||||
316
Depth-Anything-V1-main/metric_depth/zoedepth/trainers/loss.py
Normal file
316
Depth-Anything-V1-main/metric_depth/zoedepth/trainers/loss.py
Normal file
@@ -0,0 +1,316 @@
|
||||
# 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.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.cuda.amp as amp
|
||||
import numpy as np
|
||||
|
||||
|
||||
KEY_OUTPUT = 'metric_depth'
|
||||
|
||||
|
||||
def extract_key(prediction, key):
|
||||
if isinstance(prediction, dict):
|
||||
return prediction[key]
|
||||
return prediction
|
||||
|
||||
|
||||
# Main loss function used for ZoeDepth. Copy/paste from AdaBins repo (https://github.com/shariqfarooq123/AdaBins/blob/0952d91e9e762be310bb4cd055cbfe2448c0ce20/loss.py#L7)
|
||||
class SILogLoss(nn.Module):
|
||||
"""SILog loss (pixel-wise)"""
|
||||
def __init__(self, beta=0.15):
|
||||
super(SILogLoss, self).__init__()
|
||||
self.name = 'SILog'
|
||||
self.beta = beta
|
||||
|
||||
def forward(self, input, target, mask=None, interpolate=True, return_interpolated=False):
|
||||
input = extract_key(input, KEY_OUTPUT)
|
||||
if input.shape[-1] != target.shape[-1] and interpolate:
|
||||
input = nn.functional.interpolate(
|
||||
input, target.shape[-2:], mode='bilinear', align_corners=True)
|
||||
intr_input = input
|
||||
else:
|
||||
intr_input = input
|
||||
|
||||
if target.ndim == 3:
|
||||
target = target.unsqueeze(1)
|
||||
|
||||
if mask is not None:
|
||||
if mask.ndim == 3:
|
||||
mask = mask.unsqueeze(1)
|
||||
|
||||
input = input[mask]
|
||||
target = target[mask]
|
||||
|
||||
with amp.autocast(enabled=False): # amp causes NaNs in this loss function
|
||||
alpha = 1e-7
|
||||
g = torch.log(input + alpha) - torch.log(target + alpha)
|
||||
|
||||
# n, c, h, w = g.shape
|
||||
# norm = 1/(h*w)
|
||||
# Dg = norm * torch.sum(g**2) - (0.85/(norm**2)) * (torch.sum(g))**2
|
||||
|
||||
Dg = torch.var(g) + self.beta * torch.pow(torch.mean(g), 2)
|
||||
|
||||
loss = 10 * torch.sqrt(Dg)
|
||||
|
||||
if torch.isnan(loss):
|
||||
print("Nan SILog loss")
|
||||
print("input:", input.shape)
|
||||
print("target:", target.shape)
|
||||
print("G", torch.sum(torch.isnan(g)))
|
||||
print("Input min max", torch.min(input), torch.max(input))
|
||||
print("Target min max", torch.min(target), torch.max(target))
|
||||
print("Dg", torch.isnan(Dg))
|
||||
print("loss", torch.isnan(loss))
|
||||
|
||||
if not return_interpolated:
|
||||
return loss
|
||||
|
||||
return loss, intr_input
|
||||
|
||||
|
||||
def grad(x):
|
||||
# x.shape : n, c, h, w
|
||||
diff_x = x[..., 1:, 1:] - x[..., 1:, :-1]
|
||||
diff_y = x[..., 1:, 1:] - x[..., :-1, 1:]
|
||||
mag = diff_x**2 + diff_y**2
|
||||
# angle_ratio
|
||||
angle = torch.atan(diff_y / (diff_x + 1e-10))
|
||||
return mag, angle
|
||||
|
||||
|
||||
def grad_mask(mask):
|
||||
return mask[..., 1:, 1:] & mask[..., 1:, :-1] & mask[..., :-1, 1:]
|
||||
|
||||
|
||||
class GradL1Loss(nn.Module):
|
||||
"""Gradient loss"""
|
||||
def __init__(self):
|
||||
super(GradL1Loss, self).__init__()
|
||||
self.name = 'GradL1'
|
||||
|
||||
def forward(self, input, target, mask=None, interpolate=True, return_interpolated=False):
|
||||
input = extract_key(input, KEY_OUTPUT)
|
||||
if input.shape[-1] != target.shape[-1] and interpolate:
|
||||
input = nn.functional.interpolate(
|
||||
input, target.shape[-2:], mode='bilinear', align_corners=True)
|
||||
intr_input = input
|
||||
else:
|
||||
intr_input = input
|
||||
|
||||
grad_gt = grad(target)
|
||||
grad_pred = grad(input)
|
||||
mask_g = grad_mask(mask)
|
||||
|
||||
loss = nn.functional.l1_loss(grad_pred[0][mask_g], grad_gt[0][mask_g])
|
||||
loss = loss + \
|
||||
nn.functional.l1_loss(grad_pred[1][mask_g], grad_gt[1][mask_g])
|
||||
if not return_interpolated:
|
||||
return loss
|
||||
return loss, intr_input
|
||||
|
||||
|
||||
class OrdinalRegressionLoss(object):
|
||||
|
||||
def __init__(self, ord_num, beta, discretization="SID"):
|
||||
self.ord_num = ord_num
|
||||
self.beta = beta
|
||||
self.discretization = discretization
|
||||
|
||||
def _create_ord_label(self, gt):
|
||||
N,one, H, W = gt.shape
|
||||
# print("gt shape:", gt.shape)
|
||||
|
||||
ord_c0 = torch.ones(N, self.ord_num, H, W).to(gt.device)
|
||||
if self.discretization == "SID":
|
||||
label = self.ord_num * torch.log(gt) / np.log(self.beta)
|
||||
else:
|
||||
label = self.ord_num * (gt - 1.0) / (self.beta - 1.0)
|
||||
label = label.long()
|
||||
mask = torch.linspace(0, self.ord_num - 1, self.ord_num, requires_grad=False) \
|
||||
.view(1, self.ord_num, 1, 1).to(gt.device)
|
||||
mask = mask.repeat(N, 1, H, W).contiguous().long()
|
||||
mask = (mask > label)
|
||||
ord_c0[mask] = 0
|
||||
ord_c1 = 1 - ord_c0
|
||||
# implementation according to the paper.
|
||||
# ord_label = torch.ones(N, self.ord_num * 2, H, W).to(gt.device)
|
||||
# ord_label[:, 0::2, :, :] = ord_c0
|
||||
# ord_label[:, 1::2, :, :] = ord_c1
|
||||
# reimplementation for fast speed.
|
||||
ord_label = torch.cat((ord_c0, ord_c1), dim=1)
|
||||
return ord_label, mask
|
||||
|
||||
def __call__(self, prob, gt):
|
||||
"""
|
||||
:param prob: ordinal regression probability, N x 2*Ord Num x H x W, torch.Tensor
|
||||
:param gt: depth ground truth, NXHxW, torch.Tensor
|
||||
:return: loss: loss value, torch.float
|
||||
"""
|
||||
# N, C, H, W = prob.shape
|
||||
valid_mask = gt > 0.
|
||||
ord_label, mask = self._create_ord_label(gt)
|
||||
# print("prob shape: {}, ord label shape: {}".format(prob.shape, ord_label.shape))
|
||||
entropy = -prob * ord_label
|
||||
loss = torch.sum(entropy, dim=1)[valid_mask.squeeze(1)]
|
||||
return loss.mean()
|
||||
|
||||
|
||||
class DiscreteNLLLoss(nn.Module):
|
||||
"""Cross entropy loss"""
|
||||
def __init__(self, min_depth=1e-3, max_depth=10, depth_bins=64):
|
||||
super(DiscreteNLLLoss, self).__init__()
|
||||
self.name = 'CrossEntropy'
|
||||
self.ignore_index = -(depth_bins + 1)
|
||||
# self._loss_func = nn.NLLLoss(ignore_index=self.ignore_index)
|
||||
self._loss_func = nn.CrossEntropyLoss(ignore_index=self.ignore_index)
|
||||
self.min_depth = min_depth
|
||||
self.max_depth = max_depth
|
||||
self.depth_bins = depth_bins
|
||||
self.alpha = 1
|
||||
self.zeta = 1 - min_depth
|
||||
self.beta = max_depth + self.zeta
|
||||
|
||||
def quantize_depth(self, depth):
|
||||
# depth : N1HW
|
||||
# output : NCHW
|
||||
|
||||
# Quantize depth log-uniformly on [1, self.beta] into self.depth_bins bins
|
||||
depth = torch.log(depth / self.alpha) / np.log(self.beta / self.alpha)
|
||||
depth = depth * (self.depth_bins - 1)
|
||||
depth = torch.round(depth)
|
||||
depth = depth.long()
|
||||
return depth
|
||||
|
||||
|
||||
|
||||
def _dequantize_depth(self, depth):
|
||||
"""
|
||||
Inverse of quantization
|
||||
depth : NCHW -> N1HW
|
||||
"""
|
||||
# Get the center of the bin
|
||||
|
||||
|
||||
|
||||
|
||||
def forward(self, input, target, mask=None, interpolate=True, return_interpolated=False):
|
||||
input = extract_key(input, KEY_OUTPUT)
|
||||
# assert torch.all(input <= 0), "Input should be negative"
|
||||
|
||||
if input.shape[-1] != target.shape[-1] and interpolate:
|
||||
input = nn.functional.interpolate(
|
||||
input, target.shape[-2:], mode='bilinear', align_corners=True)
|
||||
intr_input = input
|
||||
else:
|
||||
intr_input = input
|
||||
|
||||
# assert torch.all(input)<=1)
|
||||
if target.ndim == 3:
|
||||
target = target.unsqueeze(1)
|
||||
|
||||
target = self.quantize_depth(target)
|
||||
if mask is not None:
|
||||
if mask.ndim == 3:
|
||||
mask = mask.unsqueeze(1)
|
||||
|
||||
# Set the mask to ignore_index
|
||||
mask = mask.long()
|
||||
input = input * mask + (1 - mask) * self.ignore_index
|
||||
target = target * mask + (1 - mask) * self.ignore_index
|
||||
|
||||
|
||||
|
||||
input = input.flatten(2) # N, nbins, H*W
|
||||
target = target.flatten(1) # N, H*W
|
||||
loss = self._loss_func(input, target)
|
||||
|
||||
if not return_interpolated:
|
||||
return loss
|
||||
return loss, intr_input
|
||||
|
||||
|
||||
|
||||
|
||||
def compute_scale_and_shift(prediction, target, mask):
|
||||
# system matrix: A = [[a_00, a_01], [a_10, a_11]]
|
||||
a_00 = torch.sum(mask * prediction * prediction, (1, 2))
|
||||
a_01 = torch.sum(mask * prediction, (1, 2))
|
||||
a_11 = torch.sum(mask, (1, 2))
|
||||
|
||||
# right hand side: b = [b_0, b_1]
|
||||
b_0 = torch.sum(mask * prediction * target, (1, 2))
|
||||
b_1 = torch.sum(mask * target, (1, 2))
|
||||
|
||||
# solution: x = A^-1 . b = [[a_11, -a_01], [-a_10, a_00]] / (a_00 * a_11 - a_01 * a_10) . b
|
||||
x_0 = torch.zeros_like(b_0)
|
||||
x_1 = torch.zeros_like(b_1)
|
||||
|
||||
det = a_00 * a_11 - a_01 * a_01
|
||||
# A needs to be a positive definite matrix.
|
||||
valid = det > 0
|
||||
|
||||
x_0[valid] = (a_11[valid] * b_0[valid] - a_01[valid] * b_1[valid]) / det[valid]
|
||||
x_1[valid] = (-a_01[valid] * b_0[valid] + a_00[valid] * b_1[valid]) / det[valid]
|
||||
|
||||
return x_0, x_1
|
||||
class ScaleAndShiftInvariantLoss(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.name = "SSILoss"
|
||||
|
||||
def forward(self, prediction, target, mask, interpolate=True, return_interpolated=False):
|
||||
|
||||
if prediction.shape[-1] != target.shape[-1] and interpolate:
|
||||
prediction = nn.functional.interpolate(prediction, target.shape[-2:], mode='bilinear', align_corners=True)
|
||||
intr_input = prediction
|
||||
else:
|
||||
intr_input = prediction
|
||||
|
||||
|
||||
prediction, target, mask = prediction.squeeze(), target.squeeze(), mask.squeeze()
|
||||
assert prediction.shape == target.shape, f"Shape mismatch: Expected same shape but got {prediction.shape} and {target.shape}."
|
||||
|
||||
scale, shift = compute_scale_and_shift(prediction, target, mask)
|
||||
|
||||
scaled_prediction = scale.view(-1, 1, 1) * prediction + shift.view(-1, 1, 1)
|
||||
|
||||
loss = nn.functional.l1_loss(scaled_prediction[mask], target[mask])
|
||||
if not return_interpolated:
|
||||
return loss
|
||||
return loss, intr_input
|
||||
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# Tests for DiscreteNLLLoss
|
||||
celoss = DiscreteNLLLoss()
|
||||
print(celoss(torch.rand(4, 64, 26, 32)*10, torch.rand(4, 1, 26, 32)*10, ))
|
||||
|
||||
d = torch.Tensor([6.59, 3.8, 10.0])
|
||||
print(celoss.dequantize_depth(celoss.quantize_depth(d)))
|
||||
@@ -0,0 +1,143 @@
|
||||
# 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 .base_trainer import BaseTrainer
|
||||
|
||||
|
||||
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.domain_classifier_loss = nn.CrossEntropyLoss()
|
||||
|
||||
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
|
||||
|
||||
Assumes all images in a batch are from the same dataset
|
||||
"""
|
||||
|
||||
images, depths_gt = batch['image'].to(
|
||||
self.device), batch['depth'].to(self.device)
|
||||
# batch['dataset'] is a tensor strings all valued either 'nyu' or 'kitti'. labels nyu -> 0, kitti -> 1
|
||||
dataset = batch['dataset'][0]
|
||||
# Convert to 0s or 1s
|
||||
domain_labels = torch.Tensor([dataset == 'kitti' for _ in range(
|
||||
images.size(0))]).to(torch.long).to(self.device)
|
||||
|
||||
# m = self.model.module if self.config.multigpu else self.model
|
||||
|
||||
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']
|
||||
domain_logits = output['domain_logits']
|
||||
|
||||
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])
|
||||
|
||||
if self.config.w_domain > 0:
|
||||
l_domain = self.domain_classifier_loss(
|
||||
domain_logits, domain_labels)
|
||||
loss = loss + self.config.w_domain * l_domain
|
||||
losses["DomainLoss"] = l_domain
|
||||
else:
|
||||
l_domain = 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 > 1 and (self.step % int(self.config.log_images_every * self.iters_per_epoch)) == 0:
|
||||
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'])
|
||||
|
||||
self.scaler.update()
|
||||
self.optimizer.zero_grad(set_to_none=True)
|
||||
|
||||
return losses
|
||||
|
||||
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]
|
||||
if 'has_valid_depth' in batch:
|
||||
if not batch['has_valid_depth']:
|
||||
return None, None
|
||||
|
||||
depths_gt = depths_gt.squeeze().unsqueeze(0).unsqueeze(0)
|
||||
with amp.autocast(enabled=self.config.use_amp):
|
||||
m = self.model.module if self.config.multigpu else self.model
|
||||
pred_depths = m(images)["metric_depth"]
|
||||
pred_depths = pred_depths.squeeze().unsqueeze(0).unsqueeze(0)
|
||||
|
||||
mask = torch.logical_and(
|
||||
depths_gt > self.config.min_depth, depths_gt < self.config.max_depth)
|
||||
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
|
||||
@@ -0,0 +1,177 @@
|
||||
# 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
|
||||
|
||||
@torch.no_grad()
|
||||
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
|
||||
|
||||
@torch.no_grad()
|
||||
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
|
||||
Reference in New Issue
Block a user