Spaces:
Runtime error
Runtime error
| """Complete Generator architecture: | |
| * OmniGenerator | |
| * Encoder | |
| * Decoders | |
| """ | |
| from pathlib import Path | |
| import traceback | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import yaml | |
| from addict import Dict | |
| from torch import softmax | |
| import climategan.strings as strings | |
| from climategan.deeplab import create_encoder, create_segmentation_decoder | |
| from climategan.depth import create_depth_decoder | |
| from climategan.masker import create_mask_decoder | |
| from climategan.painter import create_painter | |
| from climategan.tutils import init_weights, mix_noise, normalize | |
| def create_generator(opts, device="cpu", latent_shape=None, no_init=False, verbose=0): | |
| G = OmniGenerator(opts, latent_shape, verbose, no_init) | |
| if no_init: | |
| print("Sending to", device) | |
| return G.to(device) | |
| for model in G.decoders: | |
| net = G.decoders[model] | |
| if model == "s": | |
| continue | |
| if isinstance(net, nn.ModuleDict): | |
| for domain, domain_model in net.items(): | |
| init_weights( | |
| net[domain_model], | |
| init_type=opts.gen[model].init_type, | |
| init_gain=opts.gen[model].init_gain, | |
| verbose=verbose, | |
| caller=f"create_generator decoder {model} {domain}", | |
| ) | |
| else: | |
| init_weights( | |
| G.decoders[model], | |
| init_type=opts.gen[model].init_type, | |
| init_gain=opts.gen[model].init_gain, | |
| verbose=verbose, | |
| caller=f"create_generator decoder {model}", | |
| ) | |
| if G.encoder is not None and opts.gen.encoder.architecture == "base": | |
| init_weights( | |
| G.encoder, | |
| init_type=opts.gen.encoder.init_type, | |
| init_gain=opts.gen.encoder.init_gain, | |
| verbose=verbose, | |
| caller="create_generator encoder", | |
| ) | |
| print("Sending to", device) | |
| return G.to(device) | |
| class OmniGenerator(nn.Module): | |
| def __init__(self, opts, latent_shape=None, verbose=0, no_init=False): | |
| """Creates the generator. All decoders listed in opts.gen will be added | |
| to the Generator.decoders ModuleDict if opts.gen.DecoderInitial is not True. | |
| Then can be accessed as G.decoders.T or G.decoders["T"] for instance, | |
| for the image Translation decoder | |
| Args: | |
| opts (addict.Dict): configuration dict | |
| """ | |
| super().__init__() | |
| self.opts = opts | |
| self.verbose = verbose | |
| self.encoder = None | |
| if any(t in opts.tasks for t in "msd"): | |
| self.encoder = create_encoder(opts, no_init, verbose) | |
| self.decoders = {} | |
| self.painter = nn.Module() | |
| if "d" in opts.tasks: | |
| self.decoders["d"] = create_depth_decoder(opts, no_init, verbose) | |
| if self.verbose > 0: | |
| print(f" - Add {self.decoders['d'].__class__.__name__}") | |
| if "s" in opts.tasks: | |
| self.decoders["s"] = create_segmentation_decoder(opts, no_init, verbose) | |
| if "m" in opts.tasks: | |
| self.decoders["m"] = create_mask_decoder(opts, no_init, verbose) | |
| self.decoders = nn.ModuleDict(self.decoders) | |
| if "p" in self.opts.tasks: | |
| self.painter = create_painter(opts, no_init, verbose) | |
| else: | |
| if self.verbose > 0: | |
| print(" - Add Empty Painter") | |
| def __str__(self): | |
| return strings.generator(self) | |
| def encode(self, x): | |
| """ | |
| Forward x through the encoder | |
| Args: | |
| x (torch.Tensor): B3HW input tensor | |
| Returns: | |
| list: High and Low level features from the encoder | |
| """ | |
| assert self.encoder is not None | |
| return self.encoder.forward(x) | |
| def decode(self, x=None, z=None, return_z=False, return_z_depth=False): | |
| """ | |
| Comptutes the predictions of all available decoders from either x or z. | |
| If using spade for the masker with 15 channels, x *must* be provided, | |
| whether z is too or not. | |
| Args: | |
| x (torch.Tensor, optional): Input tensor (B3HW). Defaults to None. | |
| z (list, optional): List of high and low-level features as BCHW. | |
| Defaults to None. | |
| return_z (bool, optional): whether or not to return z in the dict. | |
| Defaults to False. | |
| return_z_depth (bool, optional): whether or not to return z_depth | |
| in the dict. Defaults to False. | |
| Raises: | |
| ValueError: If using spade for the masker with 15 channels but x is None | |
| Returns: | |
| dict: {task: prediction_tensor} (may include z and z_depth | |
| depending on args) | |
| """ | |
| assert x is not None or z is not None | |
| if self.opts.gen.m.use_spade and self.opts.m.spade.cond_nc == 15: | |
| if x is None: | |
| raise ValueError( | |
| "When using spade for the Masker with 15 channels," | |
| + " x MUST be provided" | |
| ) | |
| z_depth = cond = d = s = None | |
| out = {} | |
| if z is None: | |
| z = self.encode(x) | |
| if return_z: | |
| out["z"] = z | |
| if "d" in self.decoders: | |
| d, z_depth = self.decoders["d"](z) | |
| out["d"] = d | |
| if return_z_depth: | |
| out["z_depth"] = z_depth | |
| if "s" in self.decoders: | |
| s = self.decoders["s"](z, z_depth) | |
| out["s"] = s | |
| if "m" in self.decoders: | |
| if s is not None and d is not None: | |
| cond = self.make_m_cond(d, s, x) | |
| m = self.mask(z=z, cond=cond) | |
| out["m"] = m | |
| return out | |
| def sample_painter_z(self, batch_size, device, force_half=False): | |
| if self.opts.gen.p.no_z: | |
| return None | |
| z = torch.empty( | |
| batch_size, | |
| self.opts.gen.p.latent_dim, | |
| self.painter.z_h, | |
| self.painter.z_w, | |
| device=device, | |
| ).normal_(mean=0, std=1.0) | |
| if force_half: | |
| z = z.half() | |
| return z | |
| def make_m_cond(self, d, s, x=None): | |
| """ | |
| Create the masker's conditioning input when using spade from the | |
| d and s predictions and from the input x when cond_nc == 15. | |
| d and s are assumed to have the the same spatial resolution. | |
| if cond_nc == 15 then x is interpolated to match that dimension. | |
| Args: | |
| d (torch.Tensor): Raw depth prediction (B1HW) | |
| s (torch.Tensor): Raw segmentation prediction (BCHW) | |
| x (torch.Tensor, optional): Input tensor (B3hW). Mandatory | |
| when opts.gen.m.spade.cond_nc == 15 | |
| Raises: | |
| ValueError: opts.gen.m.spade.cond_nc == 15 but x is None | |
| Returns: | |
| torch.Tensor: B x cond_nc x H x W conditioning tensor. | |
| """ | |
| if self.opts.gen.m.spade.detach: | |
| d = d.detach() | |
| s = s.detach() | |
| cats = [normalize(d), softmax(s, dim=1)] | |
| if self.opts.gen.m.spade.cond_nc == 15: | |
| if x is None: | |
| raise ValueError( | |
| "When using spade for the Masker with 15 channels," | |
| + " x MUST be provided" | |
| ) | |
| cats += [ | |
| F.interpolate(x, s.shape[-2:], mode="bilinear", align_corners=True) | |
| ] | |
| return torch.cat(cats, dim=1) | |
| def mask(self, x=None, z=None, cond=None, z_depth=None, sigmoid=True): | |
| """ | |
| Create a mask from either an input x or a latent vector z. | |
| Optionally if the Masker has a spade architecture the conditioning tensor | |
| may be provided (cond). Default behavior applies an element-wise | |
| sigmoid, but can be deactivated (sigmoid=False). | |
| At least one of x or z must be provided (i.e. not None). | |
| If the Masker has a spade architecture and cond_nc == 15 then x cannot | |
| be None. | |
| Args: | |
| x (torch.Tensor, optional): Input tensor B3HW. Defaults to None. | |
| z (list, optional): High and Low level features of the encoder. | |
| Will be computed if None. Defaults to None. | |
| cond ([type], optional): [description]. Defaults to None. | |
| sigmoid (bool, optional): [description]. Defaults to True. | |
| Returns: | |
| torch.Tensor: B1HW mask tensor | |
| """ | |
| assert x is not None or z is not None | |
| if z is None: | |
| z = self.encode(x) | |
| if cond is None and self.opts.gen.m.use_spade: | |
| assert "s" in self.opts.tasks and "d" in self.opts.tasks | |
| with torch.no_grad(): | |
| d_pred, z_d = self.decoders["d"](z) | |
| s_pred = self.decoders["s"](z, z_d) | |
| cond = self.make_m_cond(d_pred, s_pred, x) | |
| if z_depth is None and self.opts.gen.m.use_dada: | |
| assert "d" in self.opts.tasks | |
| with torch.no_grad(): | |
| _, z_depth = self.decoders["d"](z) | |
| if cond is not None: | |
| device = z[0].device if isinstance(z, (tuple, list)) else z.device | |
| cond = cond.to(device) | |
| logits = self.decoders["m"](z, cond, z_depth) | |
| if not sigmoid: | |
| return logits | |
| return torch.sigmoid(logits) | |
| def paint(self, m, x, no_paste=False): | |
| """ | |
| Paints given a mask and an image | |
| calls painter(z, x * (1.0 - m)) | |
| Mask has 1s where water should be painted | |
| Args: | |
| m (torch.Tensor): Mask | |
| x (torch.Tensor): Image to paint | |
| Returns: | |
| torch.Tensor: painted image | |
| """ | |
| z_paint = self.sample_painter_z(x.shape[0], x.device) | |
| m = m.to(x.dtype) | |
| fake = self.painter(z_paint, x * (1.0 - m)) | |
| if self.opts.gen.p.paste_original_content and not no_paste: | |
| return x * (1.0 - m) + fake * m | |
| return fake | |
| def paint_cloudy(self, m, x, s, sky_idx=9, res=(8, 8), weight=0.8): | |
| """ | |
| Paints x with water in m through an intermediary cloudy image | |
| where the sky has been replaced with perlin noise to imitate clouds. | |
| The intermediary cloudy image is only used to control the painter's | |
| painting mode, probing it with a cloudy input. | |
| Args: | |
| m (torch.Tensor): water mask | |
| x (torch.Tensor): input tensor | |
| s (torch.Tensor): segmentation prediction (BCHW) | |
| sky_idx (int, optional): Index of the sky class along s's C dimension. | |
| Defaults to 9. | |
| res (tuple, optional): Perlin noise spatial resolution. Defaults to (8, 8). | |
| weight (float, optional): Intermediate image's cloud proportion | |
| (w * cloud + (1-w) * original_sky). Defaults to 0.8. | |
| Returns: | |
| torch.Tensor: painted image with original content pasted. | |
| """ | |
| sky_mask = ( | |
| torch.argmax( | |
| F.interpolate(s, x.shape[-2:], mode="bilinear"), dim=1, keepdim=True | |
| ) | |
| == sky_idx | |
| ).to(x.dtype) | |
| noised_x = mix_noise(x, sky_mask, res=res, weight=weight).to(x.dtype) | |
| fake = self.paint(m, noised_x, no_paste=True) | |
| return x * (1.0 - m) + fake * m | |
| def depth(self, x=None, z=None, return_z_depth=False): | |
| """ | |
| Compute the depth head's output | |
| Args: | |
| x (torch.Tensor, optional): Input B3HW tensor. Defaults to None. | |
| z (list, optional): High and Low level features of the encoder. | |
| Defaults to None. | |
| Returns: | |
| torch.Tensor: B1HW tensor of depth predictions | |
| """ | |
| assert x is not None or z is not None | |
| assert not (x is not None and z is not None) | |
| if z is None: | |
| z = self.encode(x) | |
| depth, z_depth = self.decoders["d"](z) | |
| if depth.shape[1] > 1: | |
| depth = torch.argmax(depth, dim=1) | |
| depth = depth / depth.max() | |
| if return_z_depth: | |
| return depth, z_depth | |
| return depth | |
| def load_val_painter(self): | |
| """ | |
| Loads a validation painter if available in opts.val.val_painter | |
| Returns: | |
| bool: operation success status | |
| """ | |
| try: | |
| # key exists in opts | |
| assert self.opts.val.val_painter | |
| # path exists | |
| ckpt_path = Path(self.opts.val.val_painter).resolve() | |
| print(ckpt_path) | |
| assert ckpt_path.exists() | |
| # path is a checkpoint path | |
| assert ckpt_path.is_file() | |
| # opts are available in that path | |
| opts_path = ckpt_path.parent.parent / "opts.yaml" | |
| assert opts_path.exists() | |
| # load opts | |
| with opts_path.open("r") as f: | |
| val_painter_opts = Dict(yaml.safe_load(f)) | |
| # load checkpoint | |
| state_dict = torch.load(ckpt_path, map_location=torch.device('cpu')) | |
| # create dummy painter from loaded opts | |
| painter = create_painter(val_painter_opts) | |
| # load state-dict in the dummy painter | |
| painter.load_state_dict( | |
| {k.replace("painter.", ""): v for k, v in state_dict["G"].items()} | |
| ) | |
| # send to current device in evaluation mode | |
| device = next(self.parameters()).device | |
| self.painter = painter.eval().to(device) | |
| # disable gradients | |
| for p in self.painter.parameters(): | |
| p.requires_grad = False | |
| # success | |
| print(" - Loaded validation-only painter") | |
| return True | |
| except Exception as e: | |
| # something happened, aborting gracefully | |
| print(traceback.format_exc()) | |
| print(e) | |
| print(">>> WARNING: error (^) in load_val_painter, aborting.") | |
| return False | |