<gh_stars>0
from training import misc
import numpy as np
import dnnlib.tflib as tflib
from matplotlib import pyplot as plt
from training import dataset


tfrecord_dir = '../../datasets/car_images_512'
label_dir = '../../datasets/car_labels/cars_v5-rxx.labels'

tflib.init_tf({'gpu_options.allow_growth': True})
dset = dataset.TFRecordDataset(tfrecord_dir, label_file=label_dir, max_label_size='full', repeat=False, buffer_mb=1, shuffle_mb=0)
tflib.init_uninitialized_vars()

x, labels = dset.get_minibatch_np(20)
zero_label = np.zeros(labels.shape)
isolated_label = np.zeros(labels.shape)
color_offset            = 1 + 67
ratio_background_offset = 1 + 67 + 12 + 18 + 10 + 8 + 5
# labels_without_background = labels
labels_without_background = np.zeros(shape=labels.shape)
# labels_without_background[:, color_offset:color_offset + 12] = np.zeros(labels.shape)[:, color_offset:color_offset + 12]
#labels_without_background = np.zeros(shape=labels.shape)

zero_label[0, 0] = 1
isolated_label[0, 0] = 1
isolated_label[:, ratio_background_offset:] = labels[:, ratio_background_offset:]
#isolated_label[0, color_offset:color_offset + 12] = labels[0, color_offset:color_offset + 12]

network_pkl = '../../results/00155-stylegan2-cars_v5_512-2gpu-config-f/network-snapshot-010467.pkl'

index = 1

_G, D_original, _Gs = misc.load_pkl(network_pkl)
D = tflib.Network('D', func_name='training.networks.networks_stylegan2_discriminator_mod_cutoff.D_stylegan2')
D.copy_vars_from(D_original)
f_maps = D.get_output_for(x, labels)
f_maps = f_maps.eval()

f_maps_zero = D.get_output_for(x, labels_without_background)
f_maps_zero = f_maps_zero.eval()

dist_fmaps = np.abs(f_maps - f_maps_zero)
dist_fmaps_mean = np.mean(dist_fmaps[index], axis=0)
fig, axs = plt.subplots(6, 6)

dist_fmaps = dist_fmaps[index]
counter = 0
for i in range(6):
    for j in range(6):
        axs[i][j].imshow(dist_fmaps[counter, :, :])
        axs[i][j].axis('off')
        counter += 1

# plt.imshow(np.mean(f_maps, axis=0))

plt.savefig('fmaps_color.png', dpi=400)
plt.show()

plt.imshow(np.transpose(x[index], [1, 2, 0]))
plt.axis('off')
plt.savefig('fmaps_input.png', dpi=300)
plt.show()

plt.imshow(dist_fmaps_mean)
plt.savefig('fmaps_diff_mean.png', dpi=300)
plt.show()

plt.close('all')
