Spaces:
Runtime error
Runtime error
| from pathlib import Path | |
| from typing import Dict, List | |
| import torchvision | |
| import torch | |
| import tops | |
| import torchvision.transforms.functional as F | |
| from .functional import hflip | |
| class RandomHorizontalFlip(torch.nn.Module): | |
| def __init__(self, p: float, flip_map=None,**kwargs): | |
| super().__init__() | |
| self.flip_ratio = p | |
| self.flip_map = flip_map | |
| if self.flip_ratio is None: | |
| self.flip_ratio = 0.5 | |
| assert 0 <= self.flip_ratio <= 1 | |
| def forward(self, container: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: | |
| if torch.rand(1) > self.flip_ratio: | |
| return container | |
| return hflip(container, self.flip_map) | |
| class CenterCrop(torch.nn.Module): | |
| """ | |
| Performs the transform on the image. | |
| NOTE: Does not transform the mask to improve runtime. | |
| """ | |
| def __init__(self, size: List[int]): | |
| super().__init__() | |
| self.size = tuple(size) | |
| def forward(self, container: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: | |
| min_size = min(container["img"].shape[1], container["img"].shape[2]) | |
| if min_size < self.size[0]: | |
| container["img"] = F.center_crop(container["img"], min_size) | |
| container["img"] = F.resize(container["img"], self.size) | |
| return container | |
| container["img"] = F.center_crop(container["img"], self.size) | |
| return container | |
| class Resize(torch.nn.Module): | |
| """ | |
| Performs the transform on the image. | |
| NOTE: Does not transform the mask to improve runtime. | |
| """ | |
| def __init__(self, size, interpolation=F.InterpolationMode.BILINEAR): | |
| super().__init__() | |
| self.size = tuple(size) | |
| self.interpolation = interpolation | |
| def forward(self, container: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: | |
| container["img"] = F.resize(container["img"], self.size, self.interpolation, antialias=True) | |
| if "semantic_mask" in container: | |
| container["semantic_mask"] = F.resize( | |
| container["semantic_mask"], self.size, F.InterpolationMode.NEAREST) | |
| if "embedding" in container: | |
| container["embedding"] = F.resize( | |
| container["embedding"], self.size, self.interpolation) | |
| if "mask" in container: | |
| container["mask"] = F.resize( | |
| container["mask"], self.size, F.InterpolationMode.NEAREST) | |
| if "E_mask" in container: | |
| container["E_mask"] = F.resize( | |
| container["E_mask"], self.size, F.InterpolationMode.NEAREST) | |
| if "maskrcnn_mask" in container: | |
| container["maskrcnn_mask"] = F.resize( | |
| container["maskrcnn_mask"], self.size, F.InterpolationMode.NEAREST) | |
| if "vertices" in container: | |
| container["vertices"] = F.resize( | |
| container["vertices"], self.size, F.InterpolationMode.NEAREST) | |
| return container | |
| def __repr__(self): | |
| repr = super().__repr__() | |
| vars_ = dict(size=self.size, interpolation=self.interpolation) | |
| return repr + " " + " ".join([f"{k}: {v}" for k, v in vars_.items()]) | |
| class InsertHRImage(torch.nn.Module): | |
| """ | |
| Resizes mask by maxpool and assumes condition is already created | |
| """ | |
| def __init__(self, size, interpolation=F.InterpolationMode.BILINEAR): | |
| super().__init__() | |
| self.size = tuple(size) | |
| self.interpolation = interpolation | |
| def forward(self, container: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: | |
| assert container["img"].dtype == torch.float32 | |
| container["img_hr"] = F.resize(container["img"], self.size, self.interpolation, antialias=True) | |
| container["condition_hr"] = F.resize(container["condition"], self.size, self.interpolation, antialias=True) | |
| mask = container["mask"] > 0 | |
| container["mask_hr"] = (torch.nn.functional.adaptive_max_pool2d(mask.logical_not().float(), output_size=self.size) > 0).logical_not().float() | |
| container["condition_hr"] = container["condition_hr"] * (1 - container["mask_hr"]) + container["img_hr"] * container["mask_hr"] | |
| return container | |
| def __repr__(self): | |
| repr = super().__repr__() | |
| vars_ = dict(size=self.size, interpolation=self.interpolation) | |
| return repr + " " | |
| class CopyHRImage(torch.nn.Module): | |
| def __init__(self) -> None: | |
| super().__init__() | |
| def forward(self, container: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: | |
| container["img_hr"] = container["img"] | |
| container["condition_hr"] = container["condition"] | |
| container["mask_hr"] = container["mask"] | |
| return container | |
| class Resize2(torch.nn.Module): | |
| """ | |
| Resizes mask by maxpool and assumes condition is already created | |
| """ | |
| def __init__(self, size, interpolation=F.InterpolationMode.BILINEAR, downsample_condition: bool = True, mask_condition= True): | |
| super().__init__() | |
| self.size = tuple(size) | |
| self.interpolation = interpolation | |
| self.downsample_condition = downsample_condition | |
| self.mask_condition = mask_condition | |
| def forward(self, container: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: | |
| # assert container["img"].dtype == torch.float32 | |
| container["img"] = F.resize(container["img"], self.size, self.interpolation, antialias=True) | |
| mask = container["mask"] > 0 | |
| container["mask"] = (torch.nn.functional.adaptive_max_pool2d(mask.logical_not().float(), output_size=self.size) > 0).logical_not().float() | |
| if self.downsample_condition: | |
| container["condition"] = F.resize(container["condition"], self.size, self.interpolation, antialias=True) | |
| if self.mask_condition: | |
| container["condition"] = container["condition"] * (1 - container["mask"]) + container["img"] * container["mask"] | |
| return container | |
| def __repr__(self): | |
| repr = super().__repr__() | |
| vars_ = dict(size=self.size, interpolation=self.interpolation) | |
| return repr + " " + " ".join([f"{k}: {v}" for k, v in vars_.items()]) | |
| class Normalize(torch.nn.Module): | |
| """ | |
| Performs the transform on the image. | |
| NOTE: Does not transform the mask to improve runtime. | |
| """ | |
| def __init__(self, mean, std, inplace, keys=["img"]): | |
| super().__init__() | |
| self.mean = mean | |
| self.std = std | |
| self.inplace = inplace | |
| self.keys = keys | |
| def forward(self, container: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: | |
| for key in self.keys: | |
| container[key] = F.normalize(container[key], self.mean, self.std, self.inplace) | |
| return container | |
| def __repr__(self): | |
| repr = super().__repr__() | |
| vars_ = dict(mean=self.mean, std=self.std, inplace=self.inplace) | |
| return repr + " " + " ".join([f"{k}: {v}" for k, v in vars_.items()]) | |
| class ToFloat(torch.nn.Module): | |
| def __init__(self, keys=["img"], norm=True) -> None: | |
| super().__init__() | |
| self.keys = keys | |
| self.gain = 255 if norm else 1 | |
| def forward(self, container: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: | |
| for key in self.keys: | |
| container[key] = container[key].float() / self.gain | |
| return container | |
| class RandomCrop(torchvision.transforms.RandomCrop): | |
| """ | |
| Performs the transform on the image. | |
| NOTE: Does not transform the mask to improve runtime. | |
| """ | |
| def forward(self, container: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: | |
| container["img"] = super().forward(container["img"]) | |
| return container | |
| class CreateCondition(torch.nn.Module): | |
| def forward(self, container: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: | |
| if container["img"].dtype == torch.uint8: | |
| container["condition"] = container["img"] * container["mask"].byte() + (1-container["mask"].byte()) * 127 | |
| return container | |
| container["condition"] = container["img"] * container["mask"] | |
| return container | |
| class CreateEmbedding(torch.nn.Module): | |
| def __init__(self, embed_path: Path, cuda=True) -> None: | |
| super().__init__() | |
| self.embed_map = torch.load(embed_path, map_location=torch.device("cpu")) | |
| if cuda: | |
| self.embed_map = tops.to_cuda(self.embed_map) | |
| def forward(self, container: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: | |
| vertices = container["vertices"] | |
| if vertices.ndim == 3: | |
| embedding = self.embed_map[vertices.long()].squeeze(dim=0) | |
| embedding = embedding.permute(2, 0, 1) * container["E_mask"] | |
| pass | |
| else: | |
| assert vertices.ndim == 4 | |
| embedding = self.embed_map[vertices.long()].squeeze(dim=1) | |
| embedding = embedding.permute(0, 3, 1, 2) * container["E_mask"] | |
| container["embedding"] = embedding | |
| container["embed_map"] = self.embed_map.clone() | |
| return container | |
| class UpdateMask(torch.nn.Module): | |
| def forward(self, container: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: | |
| container["mask"] = (container["img"] == container["condition"]).any(dim=1, keepdims=True).float() | |
| return container | |
| class LoadClassEmbedding(torch.nn.Module): | |
| def __init__(self, embedding_path: Path) -> None: | |
| super().__init__() | |
| self.embedding = torch.load(embedding_path, map_location="cpu") | |
| def forward(self, container: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: | |
| key = "_".join(container["__key__"].split("train/")[-1].split("/")[:-1]) | |
| container["class_embedding"] = self.embedding[key].view(-1) | |
| return container | |