| 173 | |
| 174 | @torch.no_grad() |
| 175 | def log_images(self, batch, only_inputs=False, log_ema=False, **kwargs): |
| 176 | log = dict() |
| 177 | x = self.get_input(batch, self.image_key) |
| 178 | x = x.to(self.device) |
| 179 | if not only_inputs: |
| 180 | xrec, posterior = self(x) |
| 181 | if x.shape[1] > 3: |
| 182 | # colorize with random projection |
| 183 | assert xrec.shape[1] > 3 |
| 184 | x = self.to_rgb(x) |
| 185 | xrec = self.to_rgb(xrec) |
| 186 | log["samples"] = self.decode(torch.randn_like(posterior.sample())) |
| 187 | log["reconstructions"] = xrec |
| 188 | if log_ema or self.use_ema: |
| 189 | with self.ema_scope(): |
| 190 | xrec_ema, posterior_ema = self(x) |
| 191 | if x.shape[1] > 3: |
| 192 | # colorize with random projection |
| 193 | assert xrec_ema.shape[1] > 3 |
| 194 | xrec_ema = self.to_rgb(xrec_ema) |
| 195 | log["samples_ema"] = self.decode(torch.randn_like(posterior_ema.sample())) |
| 196 | log["reconstructions_ema"] = xrec_ema |
| 197 | log["inputs"] = x |
| 198 | return log |
| 199 | |
| 200 | def to_rgb(self, x): |
| 201 | assert self.image_key == "segmentation" |