<jupyter_start><jupyter_text><jupyter_code>import os
import tarfile
from google.colab import drive
drive.mount('/content/drive')
from tensorflow.python.client import device_lib
device_lib.list_local_devices()
import zipfile
def un_zip(file_name,target_path):
    zip_file = zipfile.ZipFile(file_name)
    if os.path.isdir(target_path):
        pass
    else:
        os.mkdir(target_path)
    for names in zip_file.namelist():
        zip_file.extract(names,target_path)
    zip_file.close()
un_zip("/content/drive/My Drive/saliency-master.zip","/content/drive/My Drive/saliency-master")
cd /content/drive/My Drive/saliency-master/saliency-master/data/salicon
ls<jupyter_output>[0m[01;34mfixations[0m/  [01;34msaliency[0m/  [01;34mstimuli[0m/<jupyter_text>Requirements<jupyter_code>pip list |grep tensor
!pip uninstall tensorflow==1.13.1
!pip install tensorflow-gpu==1.13.1
!pip install matplotlib==3.0.3
!pip install requests==2.21.0<jupyter_output><empty_output><jupyter_text>Config<jupyter_code>import os
import sys

import numpy as np
import tensorflow as tf

PARAMS = {
    "n_epochs": 10,
    "batch_size": 1,
    "learning_rate": 1e-5,
    "device": "gpu"
}

"""The predefined input image sizes for each of the 3 datasets.
   To reproduce the results from the paper, these values should
   not be changed. They must be divisible by 8 due to the model's
   downsampling operations. Furthermore, all pretrained models
   for download were trained on these image dimensions.
"""

DIMS = {
    "image_size_salicon": (240, 320),
    "image_size_mit1003": (360, 360),
    "image_size_cat2000": (216, 384),
    "image_size_dutomron": (360, 360),
    "image_size_pascals": (360, 360),
    "image_size_osie": (240, 320),
    "image_size_fiwi": (216, 384)
}<jupyter_output><empty_output><jupyter_text>DataSet<jupyter_code>datasets_list = ["salicon", "mit1003", "cat2000",
                  "dutomron", "pascals", "osie", "fiwi"]
class SALICON:
    n_train = 10000
    n_valid = 5000

    def __init__(self, data_path):
        self._target_size = DIMS["image_size_salicon"]

        self._dir_stimuli_train = data_path + "stimuli/train"
        self._dir_stimuli_valid = data_path + "stimuli/val"

        self._dir_saliency_train = data_path + "saliency/train"
        self._dir_saliency_valid = data_path + "saliency/val"

        if not os.path.exists(data_path):
            parent_path = os.path.dirname(data_path[:-1])
            parent_path = os.path.join(parent_path, "")

            ##download.download_salicon(parent_path)

    def load_data(self):
        train_list_x = _get_file_list(self._dir_stimuli_train)
        train_list_y = _get_file_list(self._dir_saliency_train)

        _check_consistency(zip(train_list_x, train_list_y), 10000)

        train_set = _fetch_dataset((train_list_x, train_list_y),
                                   self._target_size, True)

        valid_list_x = _get_file_list(self._dir_stimuli_valid)
        valid_list_y = _get_file_list(self._dir_saliency_valid)

        _check_consistency(zip(valid_list_x, valid_list_y), 5000)

        valid_set = _fetch_dataset((valid_list_x, valid_list_y),
                                   self._target_size, False)

        return (train_set, valid_set)


class MIT1003:
    n_train = 803
    n_valid = 200

    def __init__(self, data_path):
        self._target_size =DIMS["image_size_mit1003"]

        self._dir_stimuli = data_path + "stimuli"
        self._dir_saliency = data_path + "saliency"

        if not os.path.exists(data_path):
            parent_path = os.path.dirname(data_path[:-1])
            parent_path = os.path.join(parent_path, "")

            ##download.download_mit1003(parent_path)

    def load_data(self):
        list_x = _get_file_list(self._dir_stimuli)
        list_y = _get_file_list(self._dir_saliency)

        _check_consistency(zip(list_x, list_y), 1003)

        indices = _get_random_indices(1003)
        excerpt = indices[:self.n_train]

        train_list_x = [list_x[idx] for idx in excerpt]
        train_list_y = [list_y[idx] for idx in excerpt]

        train_set = _fetch_dataset((train_list_x, train_list_y),
                                   self._target_size, True)

        excerpt = indices[self.n_train:]

        valid_list_x = [list_x[idx] for idx in excerpt]
        valid_list_y = [list_y[idx] for idx in excerpt]

        valid_set = _fetch_dataset((valid_list_x, valid_list_y),
                                   self._target_size, False)

        return (train_set, valid_set)


class CAT2000:
    n_train = 1600
    n_valid = 400

    def __init__(self, data_path):
        self._target_size =DIMS["image_size_cat2000"]

        self._dir_stimuli = data_path + "stimuli"
        self._dir_saliency = data_path + "saliency"

        if not os.path.exists(data_path):
            parent_path = os.path.dirname(data_path[:-1])
            parent_path = os.path.join(parent_path, "")

            ##download.download_cat2000(parent_path)

    def load_data(self):
        list_x = _get_file_list(self._dir_stimuli)
        list_y = _get_file_list(self._dir_saliency)

        _check_consistency(zip(list_x, list_y), 2000)

        indices = _get_random_indices(100)

        # sample uniformly from all 20 categories
        ratio = self.n_train * 100 // 2000
        excerpt = np.tile(indices[:ratio], 20)

        for idx, _ in enumerate(excerpt):
            excerpt[idx] = excerpt[idx] + idx // ratio * 100

        train_list_x = [list_x[idx] for idx in excerpt]
        train_list_y = [list_y[idx] for idx in excerpt]

        train_set = _fetch_dataset((train_list_x, train_list_y),
                                   self._target_size, True)

        # sample uniformly from all 20 categories
        ratio = self.n_valid * 100 // 2000
        excerpt = np.tile(indices[-ratio:], 20)

        for idx, _ in enumerate(excerpt):
            excerpt[idx] = excerpt[idx] + idx // ratio * 100

        valid_list_x = [list_x[idx] for idx in excerpt]
        valid_list_y = [list_y[idx] for idx in excerpt]

        valid_set = _fetch_dataset((valid_list_x, valid_list_y),
                                   self._target_size, False)

        return (train_set, valid_set)


class DUTOMRON:

    n_train = 4168
    n_valid = 1000

    def __init__(self, data_path):
        self._target_size = DIMS["image_size_dutomron"]

        self._dir_stimuli = data_path + "stimuli"
        self._dir_saliency = data_path + "saliency"

        if not os.path.exists(data_path):
            parent_path = os.path.dirname(data_path[:-1])
            parent_path = os.path.join(parent_path, "")

            ##download.download_dutomron(parent_path)

    def load_data(self):
        list_x = _get_file_list(self._dir_stimuli)
        list_y = _get_file_list(self._dir_saliency)

        _check_consistency(zip(list_x, list_y), 5168)

        indices = _get_random_indices(5168)
        excerpt = indices[:self.n_train]

        train_list_x = [list_x[idx] for idx in excerpt]
        train_list_y = [list_y[idx] for idx in excerpt]

        train_set = _fetch_dataset((train_list_x, train_list_y),
                                   self._target_size, True)

        excerpt = indices[self.n_train:]

        valid_list_x = [list_x[idx] for idx in excerpt]
        valid_list_y = [list_y[idx] for idx in excerpt]

        valid_set = _fetch_dataset((valid_list_x, valid_list_y),
                                   self._target_size, False)

        return (train_set, valid_set)


class PASCALS:
    n_train = 650
    n_valid = 200

    def __init__(self, data_path):
        self._target_size = DIMS["image_size_pascals"]

        self._dir_stimuli = data_path + "stimuli"
        self._dir_saliency = data_path + "saliency"

        if not os.path.exists(data_path):
            parent_path = os.path.dirname(data_path[:-1])
            parent_path = os.path.join(parent_path, "")

            ##download.download_pascals(parent_path)

    def load_data(self):
        list_x = _get_file_list(self._dir_stimuli)
        list_y = _get_file_list(self._dir_saliency)

        _check_consistency(zip(list_x, list_y), 850)

        indices = _get_random_indices(850)
        excerpt = indices[:self.n_train]

        train_list_x = [list_x[idx] for idx in excerpt]
        train_list_y = [list_y[idx] for idx in excerpt]

        train_set = _fetch_dataset((train_list_x, train_list_y),
                                   self._target_size, True)

        excerpt = indices[self.n_train:]

        valid_list_x = [list_x[idx] for idx in excerpt]
        valid_list_y = [list_y[idx] for idx in excerpt]

        valid_set = _fetch_dataset((valid_list_x, valid_list_y),
                                   self._target_size, False)

        return (train_set, valid_set)


class OSIE:
    n_train = 500
    n_valid = 200

    def __init__(self, data_path):
        self._target_size = DIMS["image_size_osie"]

        self._dir_stimuli = data_path + "stimuli"
        self._dir_saliency = data_path + "saliency"

        if not os.path.exists(data_path):
            parent_path = os.path.dirname(data_path[:-1])
            parent_path = os.path.join(parent_path, "")

            ##download.download_osie(parent_path)

    def load_data(self):
        list_x = _get_file_list(self._dir_stimuli)
        list_y = _get_file_list(self._dir_saliency)

        _check_consistency(zip(list_x, list_y), 700)

        indices = _get_random_indices(700)
        excerpt = indices[:self.n_train]

        train_list_x = [list_x[idx] for idx in excerpt]
        train_list_y = [list_y[idx] for idx in excerpt]

        train_set = _fetch_dataset((train_list_x, train_list_y),
                                   self._target_size, True)

        excerpt = indices[self.n_train:]

        valid_list_x = [list_x[idx] for idx in excerpt]
        valid_list_y = [list_y[idx] for idx in excerpt]

        valid_set = _fetch_dataset((valid_list_x, valid_list_y),
                                   self._target_size, False)

        return (train_set, valid_set)


class FIWI:
    n_train = 99
    n_valid = 50

    def __init__(self, data_path):
        self._target_size = DIMS["image_size_fiwi"]

        self._dir_stimuli = data_path + "stimuli"
        self._dir_saliency = data_path + "saliency"

        if not os.path.exists(data_path):
            parent_path = os.path.dirname(data_path[:-1])
            parent_path = os.path.join(parent_path, "")

            ##download.download_fiwi(parent_path)

    def load_data(self):
        list_x = _get_file_list(self._dir_stimuli)
        list_y = _get_file_list(self._dir_saliency)

        _check_consistency(zip(list_x, list_y), 149)

        indices = _get_random_indices(149)
        excerpt = indices[:self.n_train]

        train_list_x = [list_x[idx] for idx in excerpt]
        train_list_y = [list_y[idx] for idx in excerpt]

        train_set = _fetch_dataset((train_list_x, train_list_y),
                                   self._target_size, True)

        excerpt = indices[self.n_train:]

        valid_list_x = [list_x[idx] for idx in excerpt]
        valid_list_y = [list_y[idx] for idx in excerpt]

        valid_set = _fetch_dataset((valid_list_x, valid_list_y),
                                   self._target_size, False)

        return (train_set, valid_set)


class TEST:

    def __init__(self, dataset, data_path):
        #self._target_size = DIMS["image_size_%s" % dataset]

        #self._dir_stimuli_test = data_path
        self._target_size = (240, 320)
        self._dir_stimuli_test = data_path+dataset

    def load_data(self):
        test_list_x = _get_file_list(self._dir_stimuli_test)

        test_set = _fetch_dataset(test_list_x, self._target_size,
                                  False, online=True)
        return test_set<jupyter_output><empty_output><jupyter_text>Preprocessing<jupyter_code>def get_dataset_iterator(phase, dataset, data_path):
    """Entry point to make an initializable dataset iterator for either
       training or testing a model by calling the respective dataset class.

    Args:
        phase (str): Holds the current phase, which can be "train" or "test".
        dataset (str): Denotes the dataset to be used during training or the
                       suitable resizing procedure when testing a model.
        data_path (str): Points to the directory where training or testing
                         data instances are stored.

    Returns:
        iterator: An initializable dataset iterator holding the relevant data.
        initializer: An operation required to initialize the correct iterator.
    """

    if phase == "train":
        #current_module = sys.modules[__name__]
        class_name = "%s" % dataset.upper()
        if class_name=="SALICON":
            dataset_path=data_path+dataset+'/'
            print("salicon: ",dataset_path)
            dataset_class = SALICON(dataset_path)
        ###add module
        train_set, valid_set = dataset_class.load_data()
        

        iterator = tf.data.Iterator.from_structure(train_set.output_types,
                                                   train_set.output_shapes)
        next_element = iterator.get_next()

        train_init_op = iterator.make_initializer(train_set)
        valid_init_op = iterator.make_initializer(valid_set)

        return next_element, train_init_op, valid_init_op

    if phase == "test":
        data_path ="/content/drive/My Drive/"
        #test_class = TEST(dataset, data_path)
        test_class = TEST("test", data_path)
        test_set = test_class.load_data()

        iterator = tf.data.Iterator.from_structure(test_set.output_types,
                                                   test_set.output_shapes)
        next_element = iterator.get_next()

        init_op = iterator.make_initializer(test_set)

        return next_element, init_op


def postprocess_saliency_map(saliency_map, target_size):
    """This function resizes and crops a single saliency map to the original
       dimensions of the input image. The output is then encoded as a jpeg
       file suitable for saving to disk.

    Args:
        saliency_map (tensor, float32): 3D tensor that holds the values of a
                                        saliency map in the range from 0 to 1.
        target_size (tensor, int32): 1D tensor that specifies the size to which
                                     the saliency map is resized and cropped.

    Returns:
        tensor, str: A tensor of the saliency map encoded as a jpeg file.
    """

    saliency_map *= 255.0

    saliency_map = _resize_image(saliency_map, target_size, True)
    saliency_map = _crop_image(saliency_map, target_size)

    saliency_map = tf.round(saliency_map) ##四舍五入
    saliency_map = tf.cast(saliency_map, tf.uint8) ##转类型

    saliency_map_jpeg = tf.image.encode_jpeg(saliency_map, "grayscale", 100) ##编码jpeg，编码后将图片重命名写入到指定的路径下

    return saliency_map_jpeg


def _fetch_dataset(files, target_size, shuffle, online=False):
    """Here the list of file directories is shuffled (only when training),
       loaded, batched, and prefetched to ensure high GPU utilization.

    Args:
        files (list, str): A list that holds the paths to all file instances.
        target_size (tuple, int): A tuple that specifies the size to which
                                  the data will be reshaped.
        shuffle (bool): Determines whether the dataset will be shuffled or not.
        online (bool, optional): Flag that decides whether the batch size must
                                 be 1 or can take any value. Defaults to False.

    Returns:
        object: A dataset object that contains the batched and prefetched data
                instances along with their shapes and file paths.
    """

    dataset = tf.data.Dataset.from_tensor_slices(files)

    if shuffle:
        dataset = dataset.shuffle(len(files[0]))

    dataset = dataset.map(lambda *files: _parse_function(files, target_size),
                          num_parallel_calls=tf.data.experimental.AUTOTUNE)

    batch_size = 1 if online else PARAMS["batch_size"]

    dataset = dataset.batch(batch_size)
    dataset = dataset.prefetch(5)

    return dataset


def _parse_function(files, target_size):
    """This function reads image data dependent on the image type and
       whether it constitutes a stimulus or saliency map. All instances
       are then reshaped and padded to yield the target dimensionality.

    Args:
        files (tuple, str): A tuple with the paths to all file instances.
                            The first element contains the stimuli and, if
                            present, the second one the ground truth maps.
        target_size (tuple, int): A tuple that specifies the size to which
                                  the data will be reshaped.

    Returns:
        list: A list that holds the image instances along with their
              shapes and file paths.
    """

    image_list = []

    for count, filename in enumerate(files):
        image_str = tf.read_file(filename)
        channels = 3 if count == 0 else 1

        image = tf.cond(tf.image.is_jpeg(image_str),
                        lambda: tf.image.decode_jpeg(image_str,
                                                     channels=channels),
                        lambda: tf.image.decode_png(image_str,
                                                    channels=channels))
        original_size = tf.shape(image)[:2]

        image = _resize_image(image, target_size)
        image = _pad_image(image, target_size)

        image_list.append(image)

    image_list.append(original_size)
    image_list.append(files)

    return image_list


def _resize_image(image, target_size, overfull=False):
    """This resizing procedure preserves the original aspect ratio and might be
       followed by padding or cropping. Depending on whether the target size is
       smaller or larger than the current image size, the area or bicubic
       interpolation method will be utilized.

    Args:
        image (tensor, uint8): A tensor with the values of an image instance.
        target_size (tuple, int): A tuple that specifies the size to which
                                  the data will be resized.
        overfull (bool, optional): Denotes whether the resulting image will be
                                   larger or equal to the specified target
                                   size. This is crucial for the following
                                   padding or cropping. Defaults to False.

    Returns:
        tensor, float32: 4D tensor that holds the values of the resized image.

    .. seealso:: The reasoning for using either area or bicubic interpolation
                 methods is based on the OpenCV documentation recommendations.
                 [https://bit.ly/2XAavw0]
    """

    current_size = tf.shape(image)[:2]
    print(current_size,target_size)
    height_ratio = target_size[0] / current_size[0]
    width_ratio = target_size[1] / current_size[1]

    if overfull:
        target_ratio = tf.maximum(height_ratio, width_ratio)
    else:
        target_ratio = tf.minimum(height_ratio, width_ratio)

    target_size = tf.cast(current_size, tf.float64) * target_ratio
    target_size = tf.cast(tf.round(target_size), tf.int32)

    shrinking = tf.cond(tf.logical_or(current_size[0] > target_size[0],
                                      current_size[1] > target_size[1]),
                        lambda: tf.constant(True),
                        lambda: tf.constant(False))

    image = tf.expand_dims(image, 0)

    image = tf.cond(shrinking,
                    lambda: tf.image.resize_area(image, target_size,
                                                 align_corners=True),
                    lambda: tf.image.resize_bicubic(image, target_size,
                                                    align_corners=True))

    image = tf.clip_by_value(image[0], 0.0, 255.0)

    return image


def _pad_image(image, target_size):
    """A single image, either stimulus or saliency map, will be padded
       symmetrically with the constant value 126 or 0 respectively.

    Args:
        image (tensor, float32): 3D tensor with the values of the image data.
        target_size (tuple, int): A tuple that specifies the size to which
                                  the data will be resized.

    Returns:
        tensor, float32: 3D tensor that holds the values of the padded image.
    """

    current_size = tf.shape(image)

    pad_constant_value = tf.cond(tf.equal(current_size[2], 3),
                                 lambda: tf.constant(126.0),
                                 lambda: tf.constant(0.0))

    pad_vertical = (target_size[0] - current_size[0]) / 2
    pad_horizontal = (target_size[1] - current_size[1]) / 2

    pad_top = tf.floor(pad_vertical)
    pad_bottom = tf.ceil(pad_vertical)
    pad_left = tf.floor(pad_horizontal)
    pad_right = tf.ceil(pad_horizontal)

    padding = [[pad_top, pad_bottom], [pad_left, pad_right], [0, 0]]
    image = tf.pad(image, padding, constant_values=pad_constant_value)

    return image


def _crop_image(image, target_size):
    """A single saliency map will be cropped according the specified target
       size by extracting the central region of the image and correctly
       removing the added padding.

    Args:
        image (tensor, float32): 3D tensor with the values of a saliency map.
        target_size (tensor, int32): 2D tensor that specifies the size to
                                     which the data will be cropped.

    Returns:
        tensor, float32: 3D tensor that holds the values of the saliency map
                         with cropped dimensionality.
    """

    current_size = tf.shape(image)[:2]

    crop_vertical = (current_size[0] - target_size[0]) / 2
    crop_horizontal = (current_size[1] - target_size[1]) / 2

    crop_top = tf.cast(tf.floor(crop_vertical), tf.int32)
    crop_left = tf.cast(tf.floor(crop_horizontal), tf.int32)

    border_bottom = crop_top + target_size[0]
    border_right = crop_left + target_size[1]

    image = image[crop_top:border_bottom, crop_left:border_right, :]

    return image


def _get_file_list(data_path):
    """This function detects all image files within the specified parent
       directory for either training or testing. The path content cannot
       be empty, otherwise an error occurs.

    Args:
        data_path (str): Points to the directory where training or testing
                         data instances are stored.

    Returns:
        list, str: A sorted list that holds the paths to all file instances.
    """

    data_list = []

    if os.path.isfile(data_path):
        data_list.append(data_path)
    else:
        for subdir, dirs, files in os.walk(data_path):
            for file in files:
                if file.lower().endswith((".png", ".jpg", ".jpeg")):
                    data_list.append(os.path.join(subdir, file))

    data_list.sort()

    if not data_list:
        raise FileNotFoundError("No data was found")

    return data_list


def _get_random_indices(list_length):
    """A helper function to generate an array of randomly shuffled indices
       to divide the MIT1003 and CAT2000 datasets into training and validation
       instances.

    Args:
        list_length (int): The number of indices that is randomly shuffled.

    Returns:
        array, int: A 1D array that contains the shuffled data indices.
    """

    indices = np.arange(list_length)
    prng = np.random.RandomState(42)
    prng.shuffle(indices)

    return indices


def _check_consistency(zipped_file_lists, n_total_files):
    """A consistency check that makes sure all files could successfully be
       found and stimuli names correspond to the ones of ground truth maps.

    Args:
        zipped_file_lists (tuple, str): A tuple of train and valid path names.
        n_total_files (int): The total number of files expected in the list.
    """

    assert len(list(zipped_file_lists)) == n_total_files, "Files are missing"

    for file_tuple in zipped_file_lists:
        file_names = [os.path.basename(entry) for entry in list(file_tuple)]
        file_names = [os.path.splitext(entry)[0] for entry in file_names]
        file_names = [entry.replace("_fixMap", "") for entry in file_names]
        file_names = [entry.replace("_fixPts", "") for entry in file_names]

        assert len(set(file_names)) == 1, "File name mismatch"<jupyter_output><empty_output><jupyter_text>Models\Optimizers\Loss\History\ProgressBar<jupyter_code>import tensorflow as tf
from tensorflow.python.tools import freeze_graph
from tensorflow.tools.graph_transforms import TransformGraph
import os
import time
from datetime import timedelta
import matplotlib.pyplot as plt
import matplotlib.ticker as plticker
import numpy as np

class History:
    """This class represents the training history of a model. It can load the
       prior history when training continues, keeps track of the training and
       validation error, and finally plots them as a curve after each epoch.
    """

    def __init__(self, n_train_batches, n_valid_batches,
                 dataset, path, device):
        self.train_history = []
        self.valid_history = []

        self._prior_epochs = 0

        self._train_error = 0
        self._valid_error = 0

        self._n_train_batches = n_train_batches
        self._n_valid_batches = n_valid_batches

        self._path = path
        self._id = (dataset, device)

        self._get_prior_history()

    def _get_prior_history(self):
        if os.path.isfile(self._path + "train_%s_%s.txt" % self._id):
            with open(self._path + "train_%s_%s.txt" % self._id, "r") as file:
                for line in file.readlines():
                    self.train_history.append(float(line))

        if os.path.isfile(self._path + "valid_%s_%s.txt" % self._id):
            with open(self._path + "valid_%s_%s.txt" % self._id, "r") as file:
                for line in file.readlines():
                    self.valid_history.append(float(line))

        self.prior_epochs = len(self.train_history)

    def update_train_step(self, train_error):
        self._train_error += train_error

    def update_valid_step(self, valid_error):
        self._valid_error += valid_error

    def get_mean_train_error(self, reset=True):
        mean_train_error = self._train_error / self._n_train_batches

        if reset:
            self._train_error = 0

        return mean_train_error

    def get_mean_valid_error(self, reset=True):
        mean_valid_error = self._valid_error / self._n_valid_batches

        if reset:
            self._valid_error = 0

        return mean_valid_error

    def save_history(self):
        mean_train_loss = self.get_mean_train_error(False)
        mean_valid_loss = self.get_mean_valid_error(False)

        self.train_history.append(mean_train_loss)
        self.valid_history.append(mean_valid_loss)

        os.makedirs(self._path, exist_ok=True)

        with open(self._path + "train_%s_%s.txt" % self._id, "a") as file:
            file.write("%f\n" % self.train_history[-1])

        with open(self._path + "valid_%s_%s.txt" % self._id, "a") as file:
            file.write("%f\n" % self.valid_history[-1])

        if len(self.train_history) > 1:
            axes = plt.figure().gca()

            x_range = np.arange(1, len(self.train_history) + 1)

            plt.plot(x_range, self.train_history, label="train", linewidth=2)
            plt.plot(x_range, self.valid_history, label="valid", linewidth=2)

            plt.legend()
            plt.xlabel("epochs")
            plt.ylabel("error")

            locations = plticker.MultipleLocator(base=1.0)
            axes.xaxis.set_major_locator(locations)

            plt.savefig(self._path + "curve_%s_%s.png" % self._id)
            plt.close()


class Progbar:
    """This class represents a progress bar for the terminal that visualizes
       the training progress for each epoch, estimated time of accomplishment,
       and then summarizes the training and validation loss together with the
       elapsed time.
    """

    def __init__(self, n_train_data, n_train_batches,
                 batch_size, n_epochs, prior_epochs):
        self._train_time = 0
        self._valid_time = 0

        self._start_time = time.time()

        self._batch_size = batch_size

        self._n_train_data = n_train_data
        self._n_train_batches = n_train_batches

        self._target_epoch = str(n_epochs + prior_epochs).zfill(2)
        self._current_epoch = str(prior_epochs + 1).zfill(2)

    def _flush(self):
        self._train_time = 0
        self._valid_time = 0

        self._start_time = time.time()

        current_epoch_int = int(self._current_epoch) + 1
        self._current_epoch = str(current_epoch_int).zfill(2)

    def update_train_step(self, current_batch):
        current_batch += 1

        self._train_time = time.time() - self._start_time
        batch_train_time = self._train_time / current_batch

        eta = (self._n_train_batches - current_batch) * batch_train_time
        eta = str(timedelta(seconds=np.ceil(eta)))

        progress_line = "=" * (20 * current_batch // self._n_train_batches)

        current_instance = current_batch * self._batch_size
        current_instance = np.clip(current_instance, 0, self._n_train_data)

        progress_frac = "%i/%i" % (current_instance, self._n_train_data)

        information = (self._current_epoch, self._target_epoch,
                       progress_line, progress_frac, eta)

        progbar_output = "Epoch %s/%s [%-20s] %s (ETA: %s)" % information

        print(progbar_output, end="\r", flush=True)

    def update_valid_step(self):
        self._valid_time = time.time() - self._start_time - self._train_time

    def write_summary(self, mean_train_loss, mean_valid_loss):
        train_time = str(timedelta(seconds=np.ceil(self._train_time)))
        valid_time = str(timedelta(seconds=np.ceil(self._valid_time)))

        train_information = (mean_train_loss, train_time)
        valid_information = (mean_valid_loss, valid_time)

        train_output = "\n\tTrain loss: %.6f (%s)" % train_information
        valid_output = "\tValid loss: %.6f (%s)" % valid_information

        print(train_output, flush=True)
        print(valid_output, flush=True)

        self._flush()

def kld(y_true, y_pred, eps=1e-7):
    """This function computes the Kullback-Leibler divergence between ground
       truth saliency maps and their predictions. Values are first divided by
       their sum for each image to yield a distribution that adds to 1.

    Args:
        y_true (tensor, float32): A 4d tensor that holds the ground truth
                                  saliency maps with values between 0 and 255.
        y_pred (tensor, float32): A 4d tensor that holds the predicted saliency
                                  maps with values between 0 and 1.
        eps (scalar, float, optional): A small factor to avoid numerical
                                       instabilities. Defaults to 1e-7.

    Returns:
        tensor, float32: A 0D tensor that holds the averaged error.
    """

    sum_per_image = tf.reduce_sum(y_true, axis=(1, 2, 3), keep_dims=True)
    y_true /= eps + sum_per_image

    sum_per_image = tf.reduce_sum(y_pred, axis=(1, 2, 3), keep_dims=True)
    y_pred /= eps + sum_per_image

    loss = y_true * tf.log(eps + y_true / (eps + y_pred))
    loss = tf.reduce_mean(tf.reduce_sum(loss, axis=(1, 2, 3)))

    return loss

class MSINET:
    """The class representing the MSI-Net based on the VGG16 model. It
       implements a definition of the computational graph, as well as
       functions related to network training.
    """

    def __init__(self):
        self._output = None
        self._mapping = {}

        if PARAMS["device"] == "gpu":
          self._data_format = "channels_first"
          self._channel_axis = 1
          self._dims_axis = (2, 3)
        elif PARAMS["device"] == "cpu":
          self._data_format = "channels_last"
          self._channel_axis = 3
          self._dims_axis = (1, 2)

    def _encoder(self, images):
        """The encoder of the model consists of a pretrained VGG16 architecture
           with 13 convolutional layers. All dense layers are discarded and the
           last 3 layers are dilated at a rate of 2 to account for the omitted
           downsampling. Finally, the activations from 3 layers are combined.

        Args:
            images (tensor, float32): A 4D tensor that holds the RGB image
                                      batches used as input to the network.
        """

        imagenet_mean = tf.constant([103.939, 116.779, 123.68])
        imagenet_mean = tf.reshape(imagenet_mean, [1, 1, 1, 3])

        images -= imagenet_mean

        if self._data_format == "channels_first":
            images = tf.transpose(images, (0, 3, 1, 2))
        #tf.reset_default_graph()
        
        layer01 = tf.layers.conv2d(images, 64, 3,
                                  padding="same",
                                  activation=tf.nn.relu,
                                  data_format=self._data_format,
                                  name="conv1/conv1_1")

        layer02 = tf.layers.conv2d(layer01, 64, 3,
                                  padding="same",
                                  activation=tf.nn.relu,
                                  data_format=self._data_format,
                                  name="conv1/conv1_2")

        layer03 = tf.layers.max_pooling2d(layer02, 2, 2,
                                          data_format=self._data_format)

        layer04 = tf.layers.conv2d(layer03, 128, 3,
                                  padding="same",
                                  activation=tf.nn.relu,
                                  data_format=self._data_format,
                                  name="conv2/conv2_1")

        layer05 = tf.layers.conv2d(layer04, 128, 3,
                                  padding="same",
                                  activation=tf.nn.relu,
                                  data_format=self._data_format,
                                  name="conv2/conv2_2")

        layer06 = tf.layers.max_pooling2d(layer05, 2, 2,
                                          data_format=self._data_format)

        layer07 = tf.layers.conv2d(layer06, 256, 3,
                                  padding="same",
                                  activation=tf.nn.relu,
                                  data_format=self._data_format,
                                  name="conv3/conv3_1")

        layer08 = tf.layers.conv2d(layer07, 256, 3,
                                  padding="same",
                                  activation=tf.nn.relu,
                                  data_format=self._data_format,
                                  name="conv3/conv3_2")

        layer09 = tf.layers.conv2d(layer08, 256, 3,
                                  padding="same",
                                  activation=tf.nn.relu,
                                  data_format=self._data_format,
                                  name="conv3/conv3_3")

        layer10 = tf.layers.max_pooling2d(layer09, 2, 2,
                                          data_format=self._data_format)

        layer11 = tf.layers.conv2d(layer10, 512, 3,
                                  padding="same",
                                  activation=tf.nn.relu,
                                  data_format=self._data_format,
                                  name="conv4/conv4_1")

        layer12 = tf.layers.conv2d(layer11, 512, 3,
                                  padding="same",
                                  activation=tf.nn.relu,
                                  data_format=self._data_format,
                                  name="conv4/conv4_2")

        layer13 = tf.layers.conv2d(layer12, 512, 3,
                                  padding="same",
                                  activation=tf.nn.relu,
                                  data_format=self._data_format,
                                  name="conv4/conv4_3")

        layer14 = tf.layers.max_pooling2d(layer13, 2, 1,
                                          padding="same",
                                          data_format=self._data_format)

        layer15 = tf.layers.conv2d(layer14, 512, 3,
                                  padding="same",
                                  activation=tf.nn.relu,
                                  dilation_rate=2,
                                  data_format=self._data_format,
                                  name="conv5/conv5_1")

        layer16 = tf.layers.conv2d(layer15, 512, 3,
                                  padding="same",
                                  activation=tf.nn.relu,
                                  dilation_rate=2,
                                  data_format=self._data_format,
                                  name="conv5/conv5_2")

        layer17 = tf.layers.conv2d(layer16, 512, 3,
                                  padding="same",
                                  activation=tf.nn.relu,
                                  dilation_rate=2,
                                  data_format=self._data_format,
                                  name="conv5/conv5_3")

        layer18 = tf.layers.max_pooling2d(layer17, 2, 1,
                                          padding="same",
                                          data_format=self._data_format)

        encoder_output = tf.concat([layer10, layer14, layer18],
                                  axis=self._channel_axis)

        self._output = encoder_output

    def _aspp(self, features):
        """The ASPP module samples information at multiple spatial scales in
           parallel via convolutional layers with different dilation factors.
           The activations are then combined with global scene context and
           represented as a common tensor.

        Args:
            features (tensor, float32): A 4D tensor that holds the features
                                        produced by the encoder module.
        """

        branch1 = tf.layers.conv2d(features, 256, 1,
                                   padding="same",
                                   activation=tf.nn.relu,
                                   data_format=self._data_format,
                                   name="aspp/conv1_1")

        branch2 = tf.layers.conv2d(features, 256, 3,
                                   padding="same",
                                   activation=tf.nn.relu,
                                   dilation_rate=4,
                                   data_format=self._data_format,
                                   name="aspp/conv1_2")

        branch3 = tf.layers.conv2d(features, 256, 3,
                                   padding="same",
                                   activation=tf.nn.relu,
                                   dilation_rate=8,
                                   data_format=self._data_format,
                                   name="aspp/conv1_3")

        branch4 = tf.layers.conv2d(features, 256, 3,
                                   padding="same",
                                   activation=tf.nn.relu,
                                   dilation_rate=12,
                                   data_format=self._data_format,
                                   name="aspp/conv1_4")

        branch5 = tf.reduce_mean(features,
                                 axis=self._dims_axis,
                                 keepdims=True)

        branch5 = tf.layers.conv2d(branch5, 256, 1,
                                   padding="valid",
                                   activation=tf.nn.relu,
                                   data_format=self._data_format,
                                   name="aspp/conv1_5")

        shape = tf.shape(features)

        branch5 = self._upsample(branch5, shape, 1)

        context = tf.concat([branch1, branch2, branch3, branch4, branch5],
                            axis=self._channel_axis)

        aspp_output = tf.layers.conv2d(context, 256, 1,
                                       padding="same",
                                       activation=tf.nn.relu,
                                       data_format=self._data_format,
                                       name="aspp/conv2")
        self._output = aspp_output

    def _decoder(self, features):
        """The decoder model applies a series of 3 upsampling blocks that each
           performs bilinear upsampling followed by a 3x3 convolution to avoid
           checkerboard artifacts in the image space. Unlike all other layers,
           the output of the model is not modified by a ReLU.

        Args:
            features (tensor, float32): A 4D tensor that holds the features
                                        produced by the ASPP module.
        """

        shape = tf.shape(features)

        layer1 = self._upsample(features, shape, 2)

        layer2 = tf.layers.conv2d(layer1, 128, 3,
                                  padding="same",
                                  activation=tf.nn.relu,
                                  data_format=self._data_format,
                                  name="decoder/conv1")

        shape = tf.shape(layer2)

        layer3 = self._upsample(layer2, shape, 2)

        layer4 = tf.layers.conv2d(layer3, 64, 3,
                                  padding="same",
                                  activation=tf.nn.relu,
                                  data_format=self._data_format,
                                  name="decoder/conv2")

        shape = tf.shape(layer4)

        layer5 = self._upsample(layer4, shape, 2)

        layer6 = tf.layers.conv2d(layer5, 32, 3,
                                  padding="same",
                                  activation=tf.nn.relu,
                                  data_format=self._data_format,
                                  name="decoder/conv3")

        decoder_output = tf.layers.conv2d(layer6, 1, 3,
                                          padding="same",
                                          data_format=self._data_format,
                                          name="decoder/conv4")

        if self._data_format == "channels_first":
            decoder_output = tf.transpose(decoder_output, (0, 2, 3, 1))

        self._output = decoder_output

    def _upsample(self, stack, shape, factor):
        """This function resizes the input to a desired shape via the
           bilinear upsampling method.

        Args:
            stack (tensor, float32): A 4D tensor with the function input.
            shape (tensor, int32): A 1D tensor with the reference shape.
            factor (scalar, int): An integer denoting the upsampling factor.

        Returns:
            tensor, float32: A 4D tensor that holds the activations after
                             bilinear upsampling of the input.
        """

        if self._data_format == "channels_first":
            stack = tf.transpose(stack, (0, 2, 3, 1))

        stack = tf.image.resize_bilinear(stack, (shape[self._dims_axis[0]] * factor,
                                                 shape[self._dims_axis[1]] * factor))

        if self._data_format == "channels_first":
            stack = tf.transpose(stack, (0, 3, 1, 2))

        return stack

    def _normalize(self, maps, eps=1e-7):
        """This function normalizes the output values to a range
           between 0 and 1 per saliency map.

        Args:
            maps (tensor, float32): A 4D tensor that holds the model output.
            eps (scalar, float, optional): A small factor to avoid numerical
                                           instabilities. Defaults to 1e-7.
        """

        min_per_image = tf.reduce_min(maps, axis=(1, 2, 3), keep_dims=True)
        maps -= min_per_image

        max_per_image = tf.reduce_max(maps, axis=(1, 2, 3), keep_dims=True)
        maps = tf.divide(maps, eps + max_per_image, name="output")

        self._output = maps

    def _pretraining(self):
        """The first 26 variables of the model here are based on the VGG16
           network. Therefore, their names are matched to the ones of the
           pretrained VGG16 checkpoint for correct initialization.
        """

        for var in tf.global_variables()[:26]:
            key = var.name.split("/", 1)[1]
            key = key.replace("kernel:0", "weights")
            key = key.replace("bias:0", "biases")
            self._mapping[key] = var

    def forward(self, images):
        """Public method to forward RGB images through the whole network
           architecture and retrieve the resulting output.

        Args:
            images (tensor, float32): A 4D tensor that holds the values of the
                                      raw input images.

        Returns:
            tensor, float32: A 4D tensor that holds the values of the
                             predicted saliency maps.
        """

        self._encoder(images)
        self._aspp(self._output)
        self._decoder(self._output)
        self._normalize(self._output)

        return self._output

    def train(self, ground_truth, predicted_maps, learning_rate):
        """Public method to define the loss function and optimization
           algorithm for training the model.

        Args:
            ground_truth (tensor, float32): A 4D tensor with the ground truth.
            predicted_maps (tensor, float32): A 4D tensor with the predictions.
            learning_rate (scalar, float): Defines the learning rate.

        Returns:
            object: The optimizer element used to train the model.
            tensor, float32: A 0D tensor that holds the averaged error.
        """

        error =kld(ground_truth, predicted_maps)
        optimizer = tf.train.AdamOptimizer(learning_rate)
        optimizer = optimizer.minimize(error)

        return optimizer, error

    def save(self, saver, sess, dataset, path, device):
        """This saves a model checkpoint to disk and creates
           the folder if it doesn't exist yet.

        Args:
            saver (object): An object for saving the model.
            sess (object): The current TF training session.
            path (str): The path used for saving the model.
            device (str): Represents either "cpu" or "gpu".
        """

        os.makedirs(path, exist_ok=True)

        saver.save(sess, path + "model_%s_%s.ckpt" % (dataset, device),
                   write_meta_graph=False, write_state=False)

    def restore(self, sess, dataset, paths, device):
        """This function allows continued training from a prior checkpoint and
           training from scratch with the pretrained VGG16 weights. In case the
           dataset is either CAT2000 or MIT1003, a prior checkpoint based on
           the SALICON dataset is required.

        Args:
            sess (object): The current TF training session.
            dataset ([type]): The dataset used for training.
            paths (dict, str): A dictionary with all path elements.
            device (str): Represents either "cpu" or "gpu".

        Returns:
            object: A saver object for saving the model.
        """

        model_name = "model_%s_%s" % (dataset, device)
        salicon_name = "model_salicon_%s" % device
        vgg16_name = "vgg16_hybrid"

        ext1 = ".ckpt.data-00000-of-00001"
        ext2 = ".ckpt.index"

        saver = tf.train.Saver()

        if os.path.isfile(paths["latest"] + model_name + ext1) and \
           os.path.isfile(paths["latest"] + model_name + ext2):
            saver.restore(sess, paths["latest"] + model_name + ".ckpt")
        elif dataset in ("mit1003", "cat2000", "dutomron",
                         "pascals", "osie", "fiwi"):
            if os.path.isfile(paths["best"] + salicon_name + ext1) and \
               os.path.isfile(paths["best"] + salicon_name + ext2):
                saver.restore(sess, paths["best"] + salicon_name + ".ckpt")
            else:
                raise FileNotFoundError("Train model on SALICON first")
        else:
            if not (os.path.isfile(paths["weights"] + vgg16_name + ext1) or
                    os.path.isfile(paths["weights"] + vgg16_name + ext2)):
                download.download_pretrained_weights(paths["weights"],
                                                     "vgg16_hybrid")
            self._pretraining()

            loader = tf.train.Saver(self._mapping)
            loader.restore(sess, paths["weights"] + vgg16_name + ".ckpt")

        return saver

    def optimize(self, sess, dataset, path, device):
        """The best performing model is frozen, optimized for inference
           by removing unneeded training operations, and written to disk.

        Args:
            sess (object): The current TF training session.
            path (str): The path used for saving the model.
            device (str): Represents either "cpu" or "gpu".

        .. seealso:: https://bit.ly/2VBBdqQ and https://bit.ly/2W7YqBa
        """

        model_name = "model_%s_%s" % (dataset, device)
        model_path = path + model_name

        tf.train.write_graph(sess.graph.as_graph_def(),
                             path, model_name + ".pbtxt")

        freeze_graph.freeze_graph(model_path + ".pbtxt", "", False,
                                  model_path + ".ckpt", "output",
                                  "save/restore_all", "save/Const:0",
                                  model_path + ".pb", True, "")

        os.remove(model_path + ".pbtxt")

        graph_def = tf.GraphDef()

        with tf.gfile.Open(model_path + ".pb", "rb") as file:
            graph_def.ParseFromString(file.read())

        transforms = ["remove_nodes(op=Identity)",
                      "merge_duplicate_nodes",
                      "strip_unused_nodes",
                      "fold_constants(ignore_errors=true)"]

        optimized_graph_def = TransformGraph(graph_def,
                                             ["input"],
                                             ["output"],
                                             transforms)

        tf.train.write_graph(optimized_graph_def,
                             logdir=path,
                             as_text=False,
                             name=model_name + ".pb")<jupyter_output><empty_output><jupyter_text>Train<jupyter_code>import argparse
import os

import numpy as np
import tensorflow as tf
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
tf.logging.set_verbosity(tf.logging.ERROR)


def define_paths(current_path,phase,data_path):
    """A helper function to define all relevant path elements for the
       locations of data, weights, and the results from either training
       or testing a model.

    Args:
        current_path (str): The absolute path string of this script.
        args (object): A namespace object with values from command line.

    Returns:
        dict: A dictionary with all path elements.
    """


    results_path = current_path + "/results/"
    weights_path = current_path + "/weights/"

    history_path = results_path + "history/"
    images_path = results_path + "images/"
    ckpts_path = results_path + "ckpts/"
    best_path = ckpts_path + "best/"
    latest_path = ckpts_path + "latest/"

    if phase == "train":
    
        paths = {
            "data": data_path,
            "history": history_path,
            "images": images_path,
            "best": best_path,
            "latest": latest_path,
            "weights": weights_path
    }
    print(paths["data"])
    return paths

def train_model(dataset, paths, device):
    """The main function for executing network training. It loads the specified
       dataset iterator, saliency model, and helper classes. Training is then
       performed in a new session by iterating over all batches for a number of
       epochs. After validation on an independent set, the model is saved and
       the training history is updated.

    Args:
        dataset (str): Denotes the dataset to be used during training.
        paths (dict, str): A dictionary with all path elements.
        device (str): Represents either "cpu" or "gpu".
    """
    print(dataset,paths["data"])
    iterator = get_dataset_iterator("train", dataset, paths["data"])

    next_element, train_init_op, valid_init_op = iterator

    input_images, ground_truths = next_element[:2]

    input_plhd = tf.placeholder_with_default(input_images,
                                             (None, None, None, 3),
                                             name="input")
    msi_net = MSINET()

    predicted_maps = msi_net.forward(input_plhd)

    optimizer, loss = msi_net.train(ground_truths, predicted_maps,
                                    PARAMS["learning_rate"])
  
    if dataset=="salicon":
        dataset_path=paths["data"]+dataset+'/'
        dataset_class = SALICON(dataset_path)
    n_train_data = dataset_class.n_train
    n_valid_data = dataset_class.n_valid

    n_train_batches = int(np.ceil(n_train_data / PARAMS["batch_size"]))
    n_valid_batches = int(np.ceil(n_valid_data / PARAMS["batch_size"]))

    history = History(n_train_batches,
                            n_valid_batches,
                            dataset,
                            paths["history"],
                            device)

    progbar = Progbar(n_train_data,
                            n_train_batches,
                            PARAMS["batch_size"],
                            PARAMS["n_epochs"],
                            history.prior_epochs)

    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        saver = msi_net.restore(sess, dataset, paths, device)

        print(">> Start training on %s..." % dataset.upper())

        for epoch in range(PARAMS["n_epochs"]):
            print("epoch: ",epoch)
            sess.run(train_init_op)

            for batch in range(n_train_batches):
                print("batch:",batch)
                _, error = sess.run([optimizer, loss])

                history.update_train_step(error)
                progbar.update_train_step(batch)

            sess.run(valid_init_op)

            for batch in range(n_valid_batches):
                error = sess.run(loss)

                history.update_valid_step(error)
                progbar.update_valid_step()

            msi_net.save(saver, sess, dataset, paths["latest"], device)

            history.save_history()

            progbar.write_summary(history.get_mean_train_error(),
                                  history.get_mean_valid_error())

            if history.valid_history[-1] == min(history.valid_history):
                msi_net.save(saver, sess, dataset, paths["best"], device)
                msi_net.optimize(sess, dataset, paths["best"], device)

                print("\tBest model!", flush=True)



current_path="/content/drive/My Drive/saliency-master/saliency-master"
default_data_path = current_path + "/data/"
paths = define_paths(current_path, "train",default_data_path)
tf.reset_default_graph()
for dataset in datasets_list[:1]:
  train_model(dataset, paths, PARAMS["device"])<jupyter_output>[1;30;43m流式输出内容被截断，只能显示最后 5000 行内容。[0m
batch: 7623
batch: 7624
batch: 7625
batch: 7626
batch: 7627
batch: 7628
batch: 7629
batch: 7630
batch: 7631
batch: 7632
batch: 7633
batch: 7634
batch: 7635
batch: 7636
batch: 7637
batch: 7638
batch: 7639
batch: 7640
batch: 7641
batch: 7642
batch: 7643
batch: 7644
batch: 7645
batch: 7646
batch: 7647
batch: 7648
batch: 7649
batch: 7650
batch: 7651
batch: 7652
batch: 7653
batch: 7654
batch: 7655
batch: 7656
batch: 7657
batch: 7658
batch: 7659
batch: 7660
batch: 7661
batch: 7662
batch: 7663
batch: 7664
batch: 7665
batch: 7666
batch: 7667
batch: 7668
batch: 7669
batch: 7670
batch: 7671
batch: 7672
batch: 7673
batch: 7674
batch: 7675
batch: 7676
batch: 7677
batch: 7678
batch: 7679
batch: 7680
batch: 7681
batch: 7682
batch: 7683
batch: 7684
batch: 7685
batch: 7686
batch: 7687
batch: 7688
batch: 7689
batch: 7690
batch: 7691
batch: 7692
batch: 7693
batch: 7694
batch: 7695
batch: 7696
batch: 7697
batch: 7698
batch: 7699
batch: 7700
batch: 7701
batch: 7702[...]<jupyter_text>Test<jupyter_code>def test_model(dataset, paths, device):
    """The main function for executing network testing. It loads the specified
       dataset iterator and optimized saliency model. By default, when no model
       checkpoint is found locally, the pretrained weights will be downloaded.
       Testing only works for models trained on the same device as specified in
       the config file.

    Args:
        dataset (str): Denotes the dataset that was used during training.
        paths (dict, str): A dictionary with all path elements.
        device (str): Represents either "cpu" or "gpu".
    """
    
    iterator = get_dataset_iterator("test", dataset, paths["data"])

    next_element, init_op = iterator

    input_images, original_shape, file_path = next_element

    graph_def = tf.GraphDef()

    model_name = "model_%s_%s.pb" % (dataset, device)

    if os.path.isfile(paths["best"] + model_name):
        print("best")
        with tf.gfile.Open(paths["best"] + model_name, "rb") as file:
            graph_def.ParseFromString(file.read())
    else:
        print("pretrained")
        if not os.path.isfile(paths["weights"] + model_name):
            download.download_pretrained_weights(paths["weights"],
                                                 model_name[:-3])

        with tf.gfile.Open(paths["weights"] + model_name, "rb") as file:
            graph_def.ParseFromString(file.read())

    [predicted_maps] = tf.import_graph_def(graph_def,
                                           input_map={"input": input_images},
                                           return_elements=["output:0"])

    jpeg =postprocess_saliency_map(predicted_maps[0],
                                         original_shape[0])

    print(">> Start testing with %s %s model..." % (dataset.upper(), device))

    with tf.Session() as sess:
        sess.run(init_op)

        while True:
            try:
                output_file, path = sess.run([jpeg, file_path])
            except tf.errors.OutOfRangeError:
                break

            path = path[0][0].decode("utf-8")

            filename = os.path.basename(path)
            filename = os.path.splitext(filename)[0]
            filename += ".jpeg"

            os.makedirs(paths["images"], exist_ok=True)

            with open(paths["images"] + filename, "wb") as file:
                file.write(output_file)
for dataset in datasets_list[:1]:
  test_model(dataset, paths, PARAMS["device"])<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.08938831545228047, 'fraction_numerical': 0.020931848612735376, 'mean_word_length': 1.946269849196045, 'pattern_counts': {'":': 29, '<': 31, '<?xml version=': 0, '>': 38, 'https://': 3, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 222, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '389054', 'license': 'BSD-2-Clause-Patent', 'n_tokens_mistral': 18269, 'n_tokens_neox': 16634, 'n_words': 4250, 'path': 'Attention.ipynb', 'repo_name': 'Jasmine216/Project_vmaf', 'size': '59001'}