|
|
import logging |
|
|
|
|
|
from siclib.models import get_model |
|
|
from siclib.models.base_model import BaseModel |
|
|
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class GeoCalib(BaseModel): |
|
|
default_conf = { |
|
|
"backbone": {"name": "encoders.mscan"}, |
|
|
"ll_enc": {"name": "encoders.low_level_encoder"}, |
|
|
"perspective_decoder": {"name": "decoders.perspective_decoder"}, |
|
|
"optimizer": {"name": "optimization.lm_optimizer"}, |
|
|
} |
|
|
|
|
|
required_data_keys = ["image"] |
|
|
|
|
|
def _init(self, conf): |
|
|
logger.debug(f"Initializing GeoCalib with {conf}") |
|
|
self.backbone = get_model(conf.backbone["name"])(conf.backbone) |
|
|
self.ll_enc = get_model(conf.ll_enc["name"])(conf.ll_enc) if conf.ll_enc else None |
|
|
|
|
|
self.perspective_decoder = get_model(conf.perspective_decoder["name"])( |
|
|
conf.perspective_decoder |
|
|
) |
|
|
|
|
|
self.optimizer = ( |
|
|
get_model(conf.optimizer["name"])(conf.optimizer) if conf.optimizer else None |
|
|
) |
|
|
|
|
|
def _forward(self, data): |
|
|
backbone_out = self.backbone(data) |
|
|
features = {"hl": backbone_out["features"], "padding": backbone_out.get("padding", None)} |
|
|
|
|
|
if self.ll_enc is not None: |
|
|
features["ll"] = self.ll_enc(data)["features"] |
|
|
|
|
|
out = self.perspective_decoder({"features": features}) |
|
|
|
|
|
out |= { |
|
|
k: data[k] |
|
|
for k in ["image", "scales", "prior_gravity", "prior_focal", "prior_k1"] |
|
|
if k in data |
|
|
} |
|
|
|
|
|
if self.optimizer is not None: |
|
|
out |= self.optimizer(out) |
|
|
|
|
|
return out |
|
|
|
|
|
def loss(self, pred, data): |
|
|
losses, metrics = self.perspective_decoder.loss(pred, data) |
|
|
total = losses["perspective_total"] |
|
|
|
|
|
if self.optimizer is not None: |
|
|
opt_losses, param_metrics = self.optimizer.loss(pred, data) |
|
|
losses |= opt_losses |
|
|
metrics |= param_metrics |
|
|
total = total + opt_losses["param_total"] |
|
|
|
|
|
losses["total"] = total |
|
|
return losses, metrics |
|
|
|
