def sample_and_save_intermediate(model, sigmas, x=None, eps=2 * 1e-5, T=100, n_images=1, save_directory=None): """ :param model: :param sigmas: :param eps: :param T: :return: """ if not os.path.exists(save_directory): os.makedirs(save_directory) if x is None: image_size = (n_images,) + utils.get_dataset_image_size(configs.config_values.dataset) x = tf.random.uniform(shape=image_size) else: image_size = x.shape n_images = image_size[0] x_all = None for i, sigma_i in enumerate(tqdm(sigmas, desc='Sampling for each sigma')): alpha_i = eps * (sigma_i / sigmas[-1]) ** 2 idx_sigmas = tf.ones(n_images, dtype=tf.int32) * i for t in range(T): x = sample_one_step(model, x, idx_sigmas, alpha_i) if x_all is None: x_all = x else: x_all = tf.concat([x_all, x], axis=0) save_as_grid(x_all, save_directory + 'intermediate.png', rows=n_images) return x
def save_image(image, filename): mode = 'L' if image.shape[-1] == 1 else 'RGB' im = Image.new( mode, utils.get_dataset_image_size(configs.config_values.dataset)[:2]) im.paste(tf.keras.preprocessing.image.array_to_img(image)) im.save(filename + '.png', format="PNG")
def sample_and_save(model, sigmas, x=None, eps=2 * 1e-5, T=100, n_images=1, save_directory=None): """ :param model: :param sigmas: :param eps: :param T: :return: """ if not os.path.exists(save_directory): os.makedirs(save_directory) if x is None: image_size = (n_images,) + utils.get_dataset_image_size(configs.config_values.dataset) x = tf.random.uniform(shape=image_size) else: image_size = x.shape n_images = image_size[0] for i, sigma_i in enumerate(tqdm(sigmas, desc='Sampling for each sigma')): alpha_i = eps * (sigma_i / sigmas[-1]) ** 2 idx_sigmas = tf.ones(n_images, dtype=tf.int32) * i for t in range(T): x = sample_one_step(model, x, idx_sigmas, alpha_i) if (t + 1) % 10 == 0: save_as_grid(x, save_directory + f'sigma{i + 1}_t{t + 1}.png') return x
def sample_many_and_save(model, sigmas, batch_size=1000, eps=2 * 1e-5, T=100, n_images=1, save_directory=None): """ Used for sampling big amount of images (e.g. 50000) :param model: model for sampling (RefineNet) :param sigmas: sigma levels of noise :param eps: :param T: iteration per sigma level :return: Tensor of dimensions (n_images, width, height, channels) """ if not os.path.exists(save_directory): os.makedirs(save_directory) # Tuple for (n_images, width, height, channels) image_size = (n_images, ) + utils.get_dataset_image_size( configs.config_values.dataset) batch_size = min(batch_size, n_images) with tf.device("CPU"): x = tf.random.uniform(shape=image_size) x = tf.data.Dataset.from_tensor_slices(x).batch(batch_size) idx_image = 0 for i_batch, batch in enumerate( tqdm( x, total=tf.data.experimental.cardinality(x).numpy(), desc="Generating samples", )): for i, sigma_i in enumerate(sigmas): alpha_i = eps * (sigma_i / sigmas[-1])**2 idx_sigmas = tf.ones(batch.get_shape()[0], dtype=tf.int32) * i for t in range(T): batch = sample_one_step(model, batch, idx_sigmas, alpha_i) if save_directory is not None: batch = _preprocess_image_to_save(batch) for image in batch: im = Image.new("RGB", image_size[1:3]) if image_size[-1] == 1: image = tf.tile(image, [1, 1, 3]) im.paste(tf.keras.preprocessing.image.array_to_img(image)) im.save(save_directory + f"{idx_image}.png", format="PNG") idx_image += 1
def sample_many(model, sigmas, batch_size=128, eps=2 * 1e-5, T=100, n_images=1): """ Used for sampling big amount of images (e.g. 50000) :param model: model for sampling (RefineNet) :param sigmas: sigma levels of noise :param eps: :param T: iteration per sigma level :return: Tensor of dimensions (n_images, width, height, channels) """ # Tuple for (n_images, width, height, channels) image_size = (n_images, ) + utils.get_dataset_image_size( configs.config_values.dataset) batch_size = min(batch_size, n_images) with tf.device("CPU"): x = tf.random.uniform(shape=image_size) x = tf.data.Dataset.from_tensor_slices(x).batch(batch_size) x_processed = None n_processed_images = 0 for i_batch, batch in enumerate( tqdm( x, total=tf.data.experimental.cardinality(x).numpy(), desc="Generating samples", )): for i, sigma_i in enumerate(sigmas): alpha_i = eps * (sigma_i / sigmas[-1])**2 idx_sigmas = tf.ones(batch.get_shape()[0], dtype=tf.int32) * i for t in range(T): batch = sample_one_step(model, batch, idx_sigmas, alpha_i) with tf.device("CPU"): if x_processed is not None: x_processed = tf.concat([x_processed, batch], axis=0) else: x_processed = batch n_processed_images += batch_size x_processed = _preprocess_image_to_save(x_processed) return x_processed
def sample_and_save(model, sigmas, x=None, eps=2 * 1e-5, T=100, n_images=1, save_directory=None): """ :param model: :param sigmas: :param eps: :param T: :return: """ if not os.path.exists(save_directory): os.makedirs(save_directory) if x is None: image_size = [ n_images, ] + utils.get_dataset_image_size(configs.config_values.dataset) x = tf.random.uniform(shape=image_size) else: image_size = x.shape n_images = image_size[0] # # Use ground truth masks # if (configs.config_values.dataset == "masked_fashion"): # print("Using groundtruth masks...") # x = x.numpy() # fashion_test = get_train_test_data("masked_fashion")[1] # fashion_test = fashion_test.batch(n_images).take(1) # f_samples = next(iter(fashion_test)).numpy() # x[...,-1] = f_samples[...,-1] # x = tf.constant(x) for i, sigma_i in enumerate(tqdm(sigmas, desc="Sampling for each sigma")): alpha_i = eps * (sigma_i / sigmas[-1])**2 idx_sigmas = tf.ones(n_images, dtype=tf.int32) * i for t in range(T): x = sample_one_step(model, x, idx_sigmas, alpha_i) if (t + 1) % T == 0: save_as_grid(x, save_directory + f"sigma{i + 1}_t{t + 1}.png") return x
def build_model(dataset, pretrained=False): input_shape = utils.get_dataset_image_size(dataset) if configs.config_values.mask_marginals: input_shape[-1] += 1 # Append a mask channel base_model = tf.keras.applications.resnet_v2.ResNet50V2( input_shape=( input_shape[0], input_shape[1], 3, ), # if pretrained else input_shape, include_top=False if pretrained else True, weights="imagenet" if pretrained else None, classes=None if pretrained else 2, pooling="avg", classifier_activation=None, ) if pretrained: model = tf.keras.Sequential([ tf.keras.Input(shape=input_shape), tf.keras.layers.Conv2D(3, 1), # 1x1 conv to increase channels base_model, tf.keras.layers.Dense(2), ]) else: model = tf.keras.Sequential([ tf.keras.Input(shape=input_shape), tf.keras.layers.Conv2D(3, 1), # 1x1 conv to increase channels base_model, ]) logging.info(model.summary()) return model
def compute_batched_score_norms(model, x_test, masked_input=False, seed=None): # Sigma Idx -> Score score_dict = [] masks_arr = [] sigmas = utils.get_sigma_levels() input_shape = utils.get_dataset_image_size(configs.config_values.dataset) channels = input_shape[-1] progress_bar = tqdm(sigmas) for idx, sigma in enumerate(progress_bar): progress_bar.set_description("Sigma: {:.4f}".format(sigma)) _logits = [] if seed: tf.random.set_seed(seed) np.random.seed(seed) for x_batch in x_test: idx_sigmas = tf.ones(x_batch.shape[0], dtype=tf.int32) * idx score = model([x_batch, idx_sigmas]) * sigma if masked_input: _, masks = tf.split(x_batch, (channels - 1, 1), axis=-1) score = score * masks score = reduce_norm(score) _logits.append(score) score_dict.append(tf.identity(tf.concat(_logits, axis=0))) # N x L Matrix of score norms scores = tf.squeeze(tf.stack(score_dict, axis=1)) # if masked_input: # masks_arr = np.concatenate(masks_arr, axis=0) # return dict(scores=scores.numpy(), masks=masks_arr) return scores.numpy()
def build_and_train(dataset, n_epochs=25, pretrained=False): start_time = datetime.now().strftime("%y%m%d-%H%M%S") input_shape = utils.get_dataset_image_size(dataset) callbacks = [ tf.keras.callbacks.EarlyStopping( # Stop training when `val_loss` is no longer improving monitor="val_loss", # an absolute change of less than min_delta, will count as no improvement min_delta=1e-3, # "no longer improving" being defined as "for at least patience epochs" patience=20, verbose=1, ), tf.keras.callbacks.ModelCheckpoint( filepath=os.path.join(MODELDIR, "e{epoch:03d}.ckpt"), # Only save a model if `val_loss` has improved. save_best_only=True, save_weights_only=True, monitor="val_loss", verbose=1, save_freq="epoch", ), tf.keras.callbacks.ReduceLROnPlateau(monitor="val_loss", factor=0.2, min_delta=1e-3, patience=1, min_lr=1e-5), tf.keras.callbacks.TensorBoard( f"./logs/classifier/{dataset}/{start_time}", update_freq=1), ] @tf.function def remove_mask(x, l): x, _ = tf.split(x, (input_shape[-1] - 1, 1), axis=-1) return x, l model = build_model(dataset, pretrained) bce = tf.keras.losses.CategoricalCrossentropy(from_logits=True) optimizer = tf.keras.optimizers.Adam(learning_rate=3e-5) # optimizer = tf.keras.optimizers.SGD(learning_rate=0.01, momentum=0.9, nesterov=True) model.compile( optimizer=optimizer, loss=bce, metrics=[tf.keras.metrics.CategoricalAccuracy()], ) train_ds, val_ds, test_ds = load_data(dataset, include_ood=False, supervised=True) train_ds = preprocess("knee", train_ds, train=True) # train_ds.map(remove_mask, num_parallel_calls=AUTOTUNE) train_ds = train_ds.prefetch(buffer_size=AUTOTUNE) val_ds = preprocess("knee", val_ds, train=True) # val_ds.map(remove_mask, num_parallel_calls=AUTOTUNE) val_ds = val_ds.prefetch(buffer_size=AUTOTUNE) test_ds = preprocess("knee", test_ds, train=False) # test_ds.map(remove_mask, num_parallel_calls=AUTOTUNE) test_ds = test_ds.prefetch(buffer_size=AUTOTUNE) history = model.fit( train_ds, validation_data=val_ds, epochs=n_epochs, callbacks=callbacks, ) logging.info("====== Performance on Test Set ======") load_and_eval(dataset) # load_last_checkpoint(model, MODELDIR) # test_preds = model.predict(test_ds) # test_labels = np.concatenate([np.argmax(l, axis=1) for _, l in test_ds], axis=0) # test_scores = test_preds[:, 1] # inlier_test_scores = test_scores[test_labels == 0] # outlier_test_scores = test_scores[test_labels == 1] # metrics = ood_metrics( # inlier_test_scores, outlier_test_scores, verbose=True, plot=True # ) # logging.info(metrics) return history