Spaces:
Build error
Build error
File size: 8,868 Bytes
e9a2a23 4459df9 e9a2a23 4459df9 e9a2a23 4459df9 e9a2a23 4459df9 e9a2a23 4459df9 e9a2a23 4459df9 e9a2a23 5c40000 e9a2a23 4459df9 e9a2a23 4459df9 e9a2a23 4459df9 e9a2a23 4459df9 5c40000 4459df9 e9a2a23 4459df9 e9a2a23 4459df9 e9a2a23 | 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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 | """
Generate a large batch of image samples from a model and save them as a large
numpy array. This can be used to produce samples for FID evaluation.
"""
import argparse
import os
import random
import torch as th
import torch.distributed as dist
import torchvision as tv
from guided_diffusion.mask import (bbox2mask, brush_stroke_mask, get_irregular_mask, random_bbox, random_cropping_bbox)
from guided_diffusion.image_datasets import load_data
from guided_diffusion import dist_util, logger
from guided_diffusion.script_util import (
model_and_diffusion_defaults,
create_model_and_diffusion,
add_dict_to_argparser,
args_to_dict,
)
label_color_mapping = {
0: (0, 0, 0), # No information
1: (128, 64, 128), # Urban fabric
2: (244, 35, 232), # Industrial, commercial, public, military, private and transport units
3: (70, 70, 70), # Mine, dump and construction sites
4: (102, 102, 156), # Artificial non-agricultural vegetated areas
5: (190, 153, 153), # Arable land (annual crops)
6: (153, 153, 153), # Permanent crops
7: (250, 170, 30), # Pastures
8: (220, 220, 0), # Complex and mixed cultivation patterns
9: (107, 142, 35), # Orchards at the fringe of urban classes
10: (152, 251, 152), # Forests
11: (70, 130, 180), # Herbaceous vegetation associations
12: (220, 20, 60), # Open spaces with little or no vegetation
13: (255, 0, 0), # Wetlands
14: (0, 0, 142), # Water
15: (220, 220, 220) # Clouds and shadows
}
label_color_mapping_ts = th.tensor(list(label_color_mapping.values()))
print(label_color_mapping_ts)
def main():
args = create_argparser().parse_args()
args.num_classes = args.num_classes + 3 if args.inpainting else args.num_classes
dist_util.setup_dist()
logger.configure()
logger.log("creating model and diffusion...")
model, diffusion = create_model_and_diffusion(
**args_to_dict(args, model_and_diffusion_defaults().keys())
)
model.load_state_dict(
dist_util.load_state_dict(args.model_path, map_location="cpu")
)
model.to(dist_util.dev())
logger.log("creating data loader...")
data = load_data(
dataset_mode=args.dataset_mode,
data_dir=args.data_dir,
batch_size=args.batch_size,
image_size=args.image_size,
class_cond=args.class_cond,
deterministic=True,
random_crop=False,
random_flip=False,
is_train=False
)
if args.use_fp16:
model.convert_to_fp16()
model.eval()
image_path = os.path.join(args.results_path, 'images')
os.makedirs(image_path, exist_ok=True)
label_path = os.path.join(args.results_path, 'labels')
os.makedirs(label_path, exist_ok=True)
sample_path = os.path.join(args.results_path, 'samples')
os.makedirs(sample_path, exist_ok=True)
debug_path = os.path.join(args.results_path, 'debug')
os.makedirs(debug_path, exist_ok=True)
logger.log("sampling...")
all_samples = []
for i, (batch, cond) in enumerate(data):
image = ((batch + 1.0) / 2.0).cuda()
label = (cond['label_ori'].float() / 255.0).cuda()
model_kwargs = preprocess_input(batch, cond, num_classes=args.num_classes, inpainting=args.inpainting)
debug_img = model_kwargs['debug_img']
print(cond['path'])
# set hyperparameter
model_kwargs['s'] = args.s
sample_fn = (
diffusion.p_sample_loop if not args.use_ddim else diffusion.ddim_sample_loop
)
# sample = sample_fn(
# model,
# (args.batch_size, 3, image.shape[2], image.shape[3]),
# clip_denoised=args.clip_denoised,
# model_kwargs=model_kwargs,
# progress=True
# )
# sample = (sample + 1) / 2.0
sample = th.zeros(1, 3, image.shape[2], image.shape[3]).to(dist_util.dev())
gathered_samples = [th.zeros_like(sample) for _ in range(dist.get_world_size())]
dist.all_gather(gathered_samples, sample) # gather not supported with NCCL
all_samples.extend([sample.cpu().numpy() for sample in gathered_samples])
for j in range(sample.shape[0]):
print(os.path.join(image_path, cond['path'][j].split('/')[-1]))
print(os.path.join(sample_path, cond['path'][j].split('/')[-1]))
print(os.path.join(label_path, cond['path'][j].split('/')[-1]))
print(os.path.join(debug_path, cond['path'][j].split('/')[-1]))
tv.utils.save_image(image[j], os.path.join(image_path, cond['path'][j].split('/')[-1]))
tv.utils.save_image(sample[j], os.path.join(sample_path, cond['path'][j].split('/')[-1]))
tv.utils.save_image(label[j], os.path.join(label_path, cond['path'][j].split('/')[-1]))
tv.utils.save_image(debug_img[j], os.path.join(debug_path, cond['path'][j].split('/')[-1]))
logger.log(f"created {len(all_samples) * args.batch_size} samples")
# break
if len(all_samples) * args.batch_size > args.num_samples:
break
dist.barrier()
logger.log("sampling complete")
def preprocess_input(batch, data, num_classes, inpainting=False):
# move to GPU and change data types
data['label'] = data['label'].long()
# create one-hot label map
label_map = data['label']
bs, _, h, w = label_map.size()
input_label = th.FloatTensor(bs, num_classes, h, w).zero_()
debug_img = th.FloatTensor(bs, h, w, 3).zero_()
if inpainting:
unique_per_batch = [th.unique(x) for x in label_map]
bypass_label_indices = [th.randperm(len(x))[:random.randint(0, len(x))] for x in unique_per_batch]
bypass_labels = [x[bypass_label_indices[idx]] for idx, x in enumerate(unique_per_batch)]
print(unique_per_batch, bypass_labels)
for idx, bypass_label in enumerate(bypass_labels):
debug_norm_img = (batch[idx, :, :, :].permute(1, 2, 0) + 1.0) / 2.0
if random.random() < 0.5:
bypassed_label = th.where(th.isin(label_map[idx], bypass_label), 0, label_map[idx])
bypassed_image = th.where(th.isin(label_map[idx], bypass_label), batch[idx], 0)
debug_img[idx, :, :, :] = th.where(th.isin(label_map[idx], bypass_label).permute(1,2,0),
debug_norm_img, label_color_mapping_ts[label_map[idx].squeeze()]/255.0)
else:
mask = get_mask(image_size=(h, w))
bypassed_label = th.where(mask == 1, label_map[idx], 0)
bypassed_image = th.where(mask == 1, 0, batch[idx])
debug_img[idx, :, :, :] = th.where(mask.permute(1,2,0) == 1,
label_color_mapping_ts[label_map[idx].squeeze()]/255.0,
debug_norm_img)
label_map[idx, :, :, :] = bypassed_label
input_label[idx, -3:, :, :] = bypassed_image
input_semantics = input_label.scatter_(1, label_map, 1.0)
# concatenate instance map if it exists
if 'instance' in data:
inst_map = data['instance']
instance_edge_map = get_edges(inst_map)
input_semantics = th.cat((input_semantics, instance_edge_map), dim=1)
return {'y': input_semantics, 'debug_img': debug_img.permute(0, 3, 1, 2)}
def get_mask(image_size):
if random.random() < 0.5:
regular_mask = bbox2mask(image_size, random_bbox())
irregular_mask = brush_stroke_mask(image_size, )
mask = regular_mask | irregular_mask
else:
irregular_mask = brush_stroke_mask(image_size, max_loops=10)
mask = irregular_mask
return th.from_numpy(mask).permute(2,0,1)
def get_edges(t):
edge = th.ByteTensor(t.size()).zero_()
edge[:, :, :, 1:] = edge[:, :, :, 1:] | (t[:, :, :, 1:] != t[:, :, :, :-1])
edge[:, :, :, :-1] = edge[:, :, :, :-1] | (t[:, :, :, 1:] != t[:, :, :, :-1])
edge[:, :, 1:, :] = edge[:, :, 1:, :] | (t[:, :, 1:, :] != t[:, :, :-1, :])
edge[:, :, :-1, :] = edge[:, :, :-1, :] | (t[:, :, 1:, :] != t[:, :, :-1, :])
return edge.float()
def create_argparser():
defaults = dict(
data_dir="",
dataset_mode="",
clip_denoised=True,
num_samples=10000,
batch_size=1,
use_ddim=False,
model_path="",
results_path="",
is_train=False,
s=1.0,
inpainting=False
)
defaults.update(model_and_diffusion_defaults())
parser = argparse.ArgumentParser()
add_dict_to_argparser(parser, defaults)
return parser
if __name__ == "__main__":
main()
|