Update README.md

README.md

@@ -21,6 +21,145 @@ MSI-Net is a visual saliency model that predicts where humans fixate on natural
 
 <img src="https://github.com/alexanderkroner/saliency/blob/master/figures/architecture.jpg?raw=true" width="700">
 
+# Example Use
+
+### Import the dependencies
+```python
+import matplotlib.pyplot as plt
+import numpy as np
+import tensorflow as tf
+from huggingface_hub import snapshot_download
+```
+
+### Download the repo files
+```python
+hf_dir = snapshot_download(
+    repo_id="alexanderkroner/MSI-Net", allow_patterns=["*.pb", "*.jpg"]
+)
+```
+
+### Load the saliency model
+```python
+model = tf.saved_model.load(hf_dir)
+model = model.signatures["serving_default"]
+```
+
+### Load the functions for preprocessing the input and postprocessing the output
+```python
+def get_target_shape(original_shape):
+    original_aspect_ratio = original_shape[0] / original_shape[1]
+
+    square_mode = abs(original_aspect_ratio - 1.0)
+    landscape_mode = abs(original_aspect_ratio - 240 / 320)
+    portrait_mode = abs(original_aspect_ratio - 320 / 240)
+
+    best_mode = min(square_mode, landscape_mode, portrait_mode)
+
+    if best_mode == square_mode:
+        target_shape = (320, 320)
+    elif best_mode == landscape_mode:
+        target_shape = (240, 320)
+    else:
+        target_shape = (320, 240)
+
+    return target_shape
+
+
+def preprocess_input(input_image, target_shape):
+    input_tensor = tf.expand_dims(input_image, axis=0)
+
+    input_tensor = tf.image.resize(
+        input_tensor, target_shape, preserve_aspect_ratio=True
+    )
+
+    vertical_padding = target_shape[0] - input_tensor.shape[1]
+    horizontal_padding = target_shape[1] - input_tensor.shape[2]
+
+    vertical_padding_1 = vertical_padding // 2
+    vertical_padding_2 = vertical_padding - vertical_padding_1
+
+    horizontal_padding_1 = horizontal_padding // 2
+    horizontal_padding_2 = horizontal_padding - horizontal_padding_1
+
+    input_tensor = tf.pad(
+        input_tensor,
+        [
+            [0, 0],
+            [vertical_padding_1, vertical_padding_2],
+            [horizontal_padding_1, horizontal_padding_2],
+            [0, 0],
+        ],
+    )
+
+    return (
+        input_tensor,
+        [vertical_padding_1, vertical_padding_2],
+        [horizontal_padding_1, horizontal_padding_2],
+    )
+
+
+def postprocess_output(
+    output_tensor, vertical_padding, horizontal_padding, original_shape
+):
+    output_tensor = output_tensor[
+        :,
+        vertical_padding[0] : output_tensor.shape[1] - vertical_padding[1],
+        horizontal_padding[0] : output_tensor.shape[2] - horizontal_padding[1],
+        :,
+    ]
+
+    output_tensor = tf.image.resize(output_tensor, original_shape)
+
+    output_array = output_tensor.numpy().squeeze()
+    output_array = plt.cm.inferno(output_array)[..., :3]
+
+    return output_array
+```
+
+### Load and preprocess an example image
+```python
+input_image = tf.keras.utils.load_img(hf_dir + "/example.jpg")
+input_image = np.array(input_image, dtype=np.float32)
+
+original_shape = input_image.shape[:2]
+target_shape = get_target_shape(original_shape)
+
+input_tensor, vertical_padding, horizontal_padding = preprocess_input(
+    input_image, target_shape
+)
+```
+
+### Feed the input tensor to the model
+```python
+output_tensor = model(input_tensor)["output"]
+```
+
+### Postprocess and visualize the output
+```python
+saliency_map = postprocess_output(
+    output_tensor, vertical_padding, horizontal_padding, original_shape
+)
+
+alpha = 0.65
+
+blended_image = alpha * saliency_map + (1 - alpha) * input_image / 255
+
+plt.figure(figsize=(10, 5))
+
+plt.subplot(1, 2, 1)
+plt.imshow(input_image / 255)
+plt.title("Input Image")
+plt.axis("off")
+
+plt.subplot(1, 2, 2)
+plt.imshow(blended_image)
+plt.title("Saliency Map")
+plt.axis("off")
+
+plt.tight_layout()
+plt.show()
+```
+
 # Datasets
 
 Before training the model on fixation data, the encoder weights were initialized from a VGG16 backbone pre-trained on the ImageNet classification task. The model was then trained on the SALICON dataset, which consists of mouse movement recordings as a proxy for gaze measurements. Finally, the weights can be fine-tuned on human eye tracking data. MSI-Net was therefore also trained on one of the following datasets, although here we only provide the SALICON base model: