Source code for cyto_dl.nn.head.mae_head

from cyto_dl.nn.head import BaseHead


class MAEHead(BaseHead):
    def run_head(
        self,
        backbone_features,
        batch,
        stage,
        n_postprocess=1,
        run_forward=True,
        y_hat=None,
    ):
        """Run head on backbone features, calculate loss, postprocess and save image, and calculate
        metrics."""
        if run_forward:
            y_hat, mask = backbone_features
        else:
            raise ValueError("MAE head is only intended for use during training.")

        loss = (batch[self.head_name] - y_hat) ** 2
        if mask.sum() > 0:
            loss = loss[mask.bool()].mean()
        else:
            loss = loss.mean()

        return {
            "loss": loss,
            "pred": self._postprocess(y_hat, img_type="prediction", n_postprocess=n_postprocess),
            "target": self._postprocess(
                batch[self.head_name], img_type="input", n_postprocess=n_postprocess
            )
            if stage != "predict"
            else None,
            "input": self._postprocess(
                batch[self.x_key], img_type="input", n_postprocess=n_postprocess
            )
            if stage != "predict"
            else None,
        }