Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,567 Bytes
9e426da |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
import torch
from PIL import Image
from torchvision.datasets import ImageFolder
from torchvision.transforms.functional import to_tensor
from torchvision.transforms import Normalize
from src.data.dataset.metric_dataset import CenterCrop
class LocalCachedDataset(ImageFolder):
def __init__(self, root, resolution=256):
super().__init__(root)
self.transform = CenterCrop(resolution)
self.cache_root = None
def load_latent(self, latent_path):
pk_data = torch.load(latent_path)
mean = pk_data['mean'].to(torch.float32)
logvar = pk_data['logvar'].to(torch.float32)
logvar = torch.clamp(logvar, -30.0, 20.0)
std = torch.exp(0.5 * logvar)
latent = mean + torch.randn_like(mean) * std
return latent
def __getitem__(self, idx: int):
image_path, target = self.samples[idx]
latent_path = image_path.replace(self.root, self.cache_root) + ".pt"
raw_image = Image.open(image_path).convert('RGB')
raw_image = self.transform(raw_image)
raw_image = to_tensor(raw_image)
if self.cache_root is not None:
latent = self.load_latent(latent_path)
else:
latent = raw_image
return raw_image, latent, target
class ImageNet256(LocalCachedDataset):
def __init__(self, root, ):
super().__init__(root, 256)
self.cache_root = root + "_256_latent"
class ImageNet512(LocalCachedDataset):
def __init__(self, root, ):
super().__init__(root, 512)
self.cache_root = root + "_512_latent"
class PixImageNet(ImageFolder):
def __init__(self, root, resolution=256):
super().__init__(root)
self.transform = CenterCrop(resolution)
self.normalize = Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
def __getitem__(self, idx: int):
image_path, target = self.samples[idx]
raw_image = Image.open(image_path).convert('RGB')
raw_image = self.transform(raw_image)
raw_image = to_tensor(raw_image)
normalized_image = self.normalize(raw_image)
return raw_image, normalized_image, target
class PixImageNet64(PixImageNet):
def __init__(self, root, ):
super().__init__(root, 64)
class PixImageNet128(PixImageNet):
def __init__(self, root, ):
super().__init__(root, 128)
class PixImageNet256(PixImageNet):
def __init__(self, root, ):
super().__init__(root, 256)
class PixImageNet512(PixImageNet):
def __init__(self, root, ):
super().__init__(root, 512)
|